Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions .env.example
Original file line number Diff line number Diff line change
@@ -1,10 +1,17 @@
OPENAI_API_KEY=key
GOOGLE_API_KEY=key
# Agent keys
OPENAI_API_KEY=
GOOGLE_API_KEY=

# SERVICE_DESK Integration
SERVICE_DESK_URL=
SERVICE_DESK_USER=
SERVICE_DESK_TOKEN=

# Sentry configuration
SENTRY_DSN=

# Your time zone. Defaults to UTC.
TIME_ZONE=
TIME_ZONE=

# Logging level
LOG_LEVEL=
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ Lightman AI is an intelligent cybersecurity news aggregation and risk assessment
agent = "openai"
score_threshold = 8
prompt = "development"
log_level = "INFO"

[prompts]
development = "Analyze cybersecurity news for relevance to our organization."' > lightman.toml
Expand Down Expand Up @@ -129,6 +130,8 @@ Lightman AI is an intelligent cybersecurity news aggregation and risk assessment
| `--start-date` | Start date to retrieve articles | False |
| `--today` | Retrieve articles from today | False |
| `--yesterday` | Retrieve articles from yesterday | False |
| `-v` | Be more verbose on output | False |


### Environment Variables:
lightman-ai uses the following environment variables:
Expand Down
12 changes: 6 additions & 6 deletions eval/cli.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
import click
from dotenv import load_dotenv
from lightman_ai.ai.utils import AGENT_CHOICES
from lightman_ai.constants import DEFAULT_CONFIG_FILE, DEFAULT_ENV_FILE
from lightman_ai.constants import DEFAULT_AGENT, DEFAULT_CONFIG_FILE, DEFAULT_ENV_FILE, DEFAULT_SCORE
from lightman_ai.core.config import PromptConfig

from eval.classified_articles import NON_RELEVANT_ARTICLES, RELEVANT_ARTICLES
from eval.constants import DEFAULT_EVAL_CONFIG_SECTION
from eval.evaluator import eval
from eval.utils import EvalConfig, EvalFileConfig, init_eval_settings
from eval.utils import PARALLEL_WORKERS, EvalConfig, EvalFileConfig


@click.command()
Expand Down Expand Up @@ -59,14 +59,14 @@ def run(
model: str | None = None,
) -> None:
load_dotenv(env_file)
settings = init_eval_settings(env_file)

config_from_file = EvalFileConfig.get_config_from_file(config_section=config, path=config_file)
configured_prompts = PromptConfig.get_config_from_file(path=prompt_file)
eval_config = EvalConfig.init_from_dict(
{
"agent": agent or config_from_file.agent or settings.AGENT,
"agent": agent or config_from_file.agent or DEFAULT_AGENT,
"prompt": prompt or config_from_file.prompt,
"score_threshold": score or config_from_file.score_threshold or settings.SCORE,
"score_threshold": score or config_from_file.score_threshold or DEFAULT_SCORE,
"samples": samples or config_from_file.samples,
"model": model or config_from_file.model,
}
Expand All @@ -81,7 +81,7 @@ def run(
agent=eval_config.agent,
prompt=configured_prompts.get_prompt(eval_config.prompt),
model=eval_config.model,
parallel_workers=settings.PARALLEL_WORKERS,
parallel_workers=PARALLEL_WORKERS,
)


Expand Down
2 changes: 1 addition & 1 deletion eval/evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def eval(
prompt=prompt,
score=score_threshold,
logger=logger,
model=model or agent_class._default_model_name,
model=model or agent_class._DEFAULT_MODEL_NAME,
)

results_template.save()
13 changes: 1 addition & 12 deletions eval/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,23 +2,12 @@
from collections.abc import Iterable
from dataclasses import dataclass
from decimal import Decimal
from typing import Any

from lightman_ai.article.models import SelectedArticle
from lightman_ai.core.config import FileConfig, FinalConfig
from lightman_ai.core.settings import Settings
from pydantic import PositiveInt


class EvalSettings(Settings):
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)

PARALLEL_WORKERS: int = 5


def init_eval_settings(env_file: str | None = None) -> EvalSettings:
return EvalSettings(_env_file=env_file)
PARALLEL_WORKERS: int = 5


class EvalConfig(FinalConfig):
Expand Down
18 changes: 11 additions & 7 deletions src/lightman_ai/ai/base/agent.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging
from abc import ABC, abstractmethod
from typing import Never
from typing import ClassVar, Never, override

from lightman_ai.article.models import SelectedArticlesList
from pydantic_ai import Agent
Expand All @@ -9,18 +9,22 @@


class BaseAgent(ABC):
_class: type[OpenAIChatModel] | type[GoogleModel]
_default_model_name: str
_AGENT_CLASS: type[OpenAIChatModel] | type[GoogleModel]
_DEFAULT_MODEL_NAME: str
_AGENT_NAME: ClassVar[str]

def __init__(self, system_prompt: str, model: str | None = None, logger: logging.Logger | None = None) -> None:
agent_model = self._class(model or self._default_model_name)
selected_model = model or self._DEFAULT_MODEL_NAME
agent_model = self._AGENT_CLASS(selected_model)
self.agent: Agent[Never, SelectedArticlesList] = Agent(
model=agent_model, output_type=SelectedArticlesList, system_prompt=system_prompt
)
self.logger = logger or logging.getLogger("lightman")
self.logger.info("Selected %s's %s model", self, selected_model)

def get_prompt_result(self, prompt: str) -> SelectedArticlesList:
return self._run_prompt(prompt)
@override
def __str__(self) -> str:
return self._AGENT_NAME

@abstractmethod
def _run_prompt(self, prompt: str) -> SelectedArticlesList: ...
def run_prompt(self, prompt: str) -> SelectedArticlesList: ...
7 changes: 4 additions & 3 deletions src/lightman_ai/ai/gemini/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,12 @@
class GeminiAgent(BaseAgent):
"""Class that provides an interface to operate with the Gemini model."""

_class = GoogleModel
_default_model_name = "gemini-2.5-pro"
_AGENT_CLASS = GoogleModel
_DEFAULT_MODEL_NAME = "gemini-2.5-pro"
_AGENT_NAME = "Gemini"

@override
def _run_prompt(self, prompt: str) -> SelectedArticlesList:
def run_prompt(self, prompt: str) -> SelectedArticlesList:
with map_gemini_exceptions():
result = self.agent.run_sync(prompt)
return result.output
7 changes: 4 additions & 3 deletions src/lightman_ai/ai/openai/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,16 @@
class OpenAIAgent(BaseAgent):
"""Class that provides an interface to operate with the OpenAI model."""

_class = OpenAIChatModel
_default_model_name = "gpt-4.1"
_AGENT_CLASS = OpenAIChatModel
_DEFAULT_MODEL_NAME = "gpt-4.1"
_AGENT_NAME = "OpenAI"

def _execute_agent(self, prompt: str) -> AgentRunResult[SelectedArticlesList]:
with map_openai_exceptions():
return self.agent.run_sync(prompt)

@override
def _run_prompt(self, prompt: str) -> SelectedArticlesList:
def run_prompt(self, prompt: str) -> SelectedArticlesList:
try:
result = self._execute_agent(prompt)
except LimitTokensExceededError as err:
Expand Down
48 changes: 40 additions & 8 deletions src/lightman_ai/cli.py
Original file line number Diff line number Diff line change
@@ -1,20 +1,33 @@
import logging
import os
from datetime import date
from importlib import metadata

import click
from dotenv import load_dotenv
from lightman_ai.ai.utils import AGENT_CHOICES
from lightman_ai.constants import DEFAULT_CONFIG_FILE, DEFAULT_CONFIG_SECTION, DEFAULT_ENV_FILE
from lightman_ai.constants import (
DEFAULT_AGENT,
DEFAULT_CONFIG_FILE,
DEFAULT_CONFIG_SECTION,
DEFAULT_ENV_FILE,
DEFAULT_LOG_LEVEL,
DEFAULT_SCORE,
DEFAULT_TIME_ZONE,
VERBOSE_LOG_LEVEL,
)
from lightman_ai.core.config import FileConfig, FinalConfig, PromptConfig
from lightman_ai.core.exceptions import ConfigNotFoundError, InvalidConfigError, PromptNotFoundError
from lightman_ai.core.sentry import configure_sentry
from lightman_ai.core.settings import Settings
from lightman_ai.exceptions import MultipleDateSourcesError
from lightman_ai.main import lightman
from lightman_ai.utils import get_start_date

logger = logging.getLogger("lightman")
logging.basicConfig(
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)


def get_version() -> str:
Expand Down Expand Up @@ -72,6 +85,7 @@ def entry_point() -> None:
@click.option("--start-date", type=click.DateTime(formats=["%Y-%m-%d"]), help="Start date to retrieve articles")
@click.option("--today", is_flag=True, help="Retrieve articles from today.")
@click.option("--yesterday", is_flag=True, help="Retrieve articles from yesterday.")
@click.option("-v", is_flag=True, help="Be more verbose on output.")
def run(
agent: str,
prompt: str,
Expand All @@ -85,18 +99,30 @@ def run(
start_date: date | None,
today: bool,
yesterday: bool,
v: bool,
) -> int:
"""
Entrypoint of the application.

Holds no logic. It loads the configuration, calls the main method and returns 0 when succesful .
"""
load_dotenv(env_file or DEFAULT_ENV_FILE)
configure_sentry()

settings = Settings.try_load_from_file(env_file)
if v:
logger.setLevel(VERBOSE_LOG_LEVEL)
else:
try:
env_log_level = os.getenv("LOG_LEVEL")
log_level = env_log_level.upper() if env_log_level else DEFAULT_LOG_LEVEL
logger.setLevel(log_level)
except ValueError:
logger.setLevel(DEFAULT_LOG_LEVEL)
logger.warning("Invalid logging level. Using default value.")

configure_sentry(logger.level)

try:
start_datetime = get_start_date(settings, yesterday, today, start_date)
start_datetime = get_start_date(os.getenv("TIME_ZONE", DEFAULT_TIME_ZONE), yesterday, today, start_date)
except MultipleDateSourcesError as e:
raise click.UsageError(e.args[0]) from e

Expand All @@ -105,9 +131,9 @@ def run(
config_from_file = FileConfig.get_config_from_file(config_section=config, path=config_file)
final_config = FinalConfig.init_from_dict(
data={
"agent": agent or config_from_file.agent or settings.AGENT,
"agent": agent or config_from_file.agent or DEFAULT_AGENT,
"prompt": prompt or config_from_file.prompt,
"score_threshold": score or config_from_file.score_threshold or settings.SCORE,
"score_threshold": score or config_from_file.score_threshold or DEFAULT_SCORE,
"model": model or config_from_file.model,
}
)
Expand All @@ -126,5 +152,11 @@ def run(
start_date=start_datetime,
)
relevant_articles_metadata = [f"{article.title} ({article.link})" for article in relevant_articles]
logger.warning("Found these articles: \n- %s", "\n- ".join(relevant_articles_metadata))

if relevant_articles_metadata:
articles = f"Found these articles:\n* {'\n* '.join(relevant_articles_metadata)} "
click.echo(click.style(articles))
else:
click.echo(click.style("No relevant articles found."))

return 0
7 changes: 7 additions & 0 deletions src/lightman_ai/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,10 @@
DEFAULT_CONFIG_FILE = "lightman.toml"

DEFAULT_ENV_FILE = ".env"

DEFAULT_LOG_LEVEL = "WARNING"
VERBOSE_LOG_LEVEL = "INFO"

DEFAULT_AGENT = "openai"
DEFAULT_SCORE = 8
DEFAULT_TIME_ZONE = "UTC"
27 changes: 10 additions & 17 deletions src/lightman_ai/core/sentry.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,31 +5,24 @@
logger = logging.getLogger("lightman")


def configure_sentry() -> None:
def configure_sentry(log_level: int) -> None:
"""Configure Sentry for error tracking."""
try:
import sentry_sdk # noqa: PLC0415
from sentry_sdk.integrations.logging import LoggingIntegration # noqa: PLC0415
except ImportError:
logger.warning(
"Could not initialize sentry, it is not installed! Add it by installing the project with `lightman-ai[sentry]`."
)
if os.getenv("SENTRY_DSN"):
logger.warning(
"Could not initialize sentry, it is not installed! Add it by installing the project with `lightman-ai[sentry]`."
)
return

try:
if not os.getenv("SENTRY_DSN"):
logger.info("SENTRY_DSN not configured, skipping Sentry initialization")
return

logging_level_str = os.getenv("LOGGING_LEVEL", "ERROR").upper()
try:
logging_level = getattr(logging, logging_level_str, logging.ERROR)
except AttributeError:
logger.warning("The specified logging level `%s` does not exist. Defaulting to ERROR.", logging_level_str)
logging_level = logging.ERROR
if not os.getenv("SENTRY_DSN"):
logger.warning("SENTRY_DSN not configured, skipping Sentry initialization")
return

# Set up logging integration
sentry_logging = LoggingIntegration(level=logging.INFO, event_level=logging_level)
try:
sentry_logging = LoggingIntegration(level=logging.INFO, event_level=log_level)

sentry_sdk.init(
release=metadata.version("lightman-ai"),
Expand Down
30 changes: 0 additions & 30 deletions src/lightman_ai/core/settings.py

This file was deleted.

Loading