diff --git a/.env-template b/.env-template index 0e09a39..446720a 100644 --- a/.env-template +++ b/.env-template @@ -1,5 +1,5 @@ -LLM_FAMILY="gemini" or "claude" or "claude-aws-bedrock" or "nova-pro-aws-bedrock" -GEMINI_API_KEY="your key if using gemini" +LLM_FAMILY="gemini" or "claude" or "claude-aws-bedrock" or "nova-lite-aws-bedrock" +GOOGLE_API_KEY="your key if using gemini" ANTHROPIC_API_KEY="your key if using Claude" AWS_ACCESS_KEY_ID="your aws access key id if using Claude via aws bedrock" AWS_SECRET_ACCESS_KEY_ID="your aws secret access key id if using Claude via aws bedrock" diff --git a/CLAUDE.md b/CLAUDE.md index 0fc7269..11754c9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -58,16 +58,16 @@ pre-commit run --all-files # Run pre-commit hooks ### Model Support - **Gemini 2.0 Flash** (default, preferred for cost and context window) -- **Claude 3.5 Haiku** (experimental, rate limited) +- **Claude 4.5 Haiku** (experimental, rate limited) - **Claude via AWS Bedrock** (experimental) -- **Nova Pro via AWS Bedrock** (experimental) +- **Nova 2 Lite via AWS Bedrock** (experimental) ## Configuration ### Environment Variables Required in `.env` file (copy from `.env-template`): -- `LLM_FAMILY`: "gemini", "claude", "claude-aws-bedrock", or "nova-pro-aws-bedrock" -- `GEMINI_API_KEY`: For Gemini models +- `LLM_FAMILY`: "gemini", "claude", "claude-aws-bedrock", or "nova-lite-aws-bedrock" +- `GOOGLE_API_KEY`: For Gemini models - `ANTHROPIC_API_KEY`: For Claude models - `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_REGION`: For AWS Bedrock diff --git a/README.md b/README.md index 4dee8c3..1110f7b 100644 --- a/README.md +++ b/README.md @@ -23,9 +23,9 @@ As for everything that uses LLM, **check your outputs** - it can make mistakes. As its underlying LLM, you can choose to run it either with: * Gemini (it will use 2.0 Flash) [preferred and default option] -* Claude (it will use Haiku 3.5) [experimental] -* Claude via AWS Bedrock (it will use Haiku 3.5) [experimental] -* Nova Pro via AWS Bedrock [experimental] +* Claude (it will use Haiku 4.5) [experimental] +* Claude via AWS Bedrock (it will use Haiku 4.5) [experimental] +* Nova 2 Lite via AWS Bedrock [experimental] Gemini is preferred because: * it has a free tier - we privilege cost-effectiveness over speed, which means for short conversations you should be within the quotas of the free tier diff --git a/evals/evals.py b/evals/evals.py index 5d60a18..cf1a9b0 100644 --- a/evals/evals.py +++ b/evals/evals.py @@ -2,6 +2,7 @@ import os import boto3 +from botocore.exceptions import ClientError, NoCredentialsError from dotenv import load_dotenv from evals_allower import run_evals as run_evals_allower from evals_article import run_evals as run_evals_article @@ -22,13 +23,13 @@ async def main(): "gemini", "claude", "claude-aws-bedrock", - "nova-pro-aws-bedrock", + "nova-lite-aws-bedrock", ]: - if model_family == "gemini" and os.getenv("GEMINI_API_KEY") is None: + if model_family == "gemini" and os.getenv("GOOGLE_API_KEY") is None: console.print( """ - [bold red]GEMINI_API_KEY environment variable is not set. + [bold red]GOOGLE_API_KEY environment variable is not set. Skipping evals.[/bold red] """ ) @@ -43,13 +44,13 @@ async def main(): ) continue - if model_family == "claude-aws-bedrock": + if model_family in ("claude-aws-bedrock", "nova-lite-aws-bedrock"): try: _ = boto3.client("sts").get_caller_identity() - except boto3.exceptions.ClientError: + except (ClientError, NoCredentialsError): console.print( - """[bold red]AWS credentials are not set or invalid. - Skipping CLAUDE AWS Bedrock evals.[/bold red]""" + f"""[bold red]AWS credentials are not set or invalid. + Skipping {model_family} evals.[/bold red]""" ) continue diff --git a/pyproject.toml b/pyproject.toml index f1409b4..a68d6b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,11 +7,11 @@ build-backend = "setuptools.build_meta" [project] name = "askademic" -version = "2.0.0" +version = "2.1.0" dependencies = [ - "pydantic==2.10.6", - "pydantic-ai==0.2.15", - "logfire==3.18.0", + "pydantic==2.11.7", + "pydantic-ai==1.44.0", + "logfire==4.16.0", "requests==2.32.3", "pandas==2.2.3", "tabulate==0.9.0", @@ -44,3 +44,7 @@ askademic = "askademic.main:main" [tool.pytest.ini_options] pythonpath = ["src"] +asyncio_default_fixture_loop_scope = "function" +filterwarnings = [ + "ignore:builtin type Swig:DeprecationWarning", +] diff --git a/src/askademic/article.py b/src/askademic/article.py index c115247..61582fb 100644 --- a/src/askademic/article.py +++ b/src/askademic/article.py @@ -3,7 +3,6 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent -from pydantic_ai.models import Model from pydantic_ai.settings import ModelSettings from askademic.prompts.general import ( @@ -61,7 +60,9 @@ class ArticleRetrievalResponse(BaseModel): class ArticleAgent: - def __init__(self, model: Model, model_settings: ModelSettings = None, use_cache: bool = True): + def __init__( + self, model: str, model_settings: ModelSettings = None, use_cache: bool = True + ): self._get_article = get_article self.use_cache = use_cache diff --git a/src/askademic/constants.py b/src/askademic/constants.py index bb6f7d3..887e81f 100644 --- a/src/askademic/constants.py +++ b/src/askademic/constants.py @@ -7,11 +7,13 @@ - Type "help" to see instructions again """ -# LLM model IDs -GEMINI_2_FLASH_MODEL_ID = "gemini-2.0-flash" -CLAUDE_HAIKU_3_5_MODEL_ID = "anthropic:claude-3-5-haiku-latest" -CLAUDE_HAIKU_3_5_BEDROCK_MODEL_ID = "{region}.anthropic.claude-3-5-haiku-20241022-v1:0" -NOVA_PRO_BEDROCK_MODEL_ID = "{region}.amazon.nova-pro-v1:0" +# LLM model IDs (using provider:model syntax) +GEMINI_2_FLASH_MODEL_ID = "google-gla:gemini-2.0-flash" +CLAUDE_HAIKU_4_5_MODEL_ID = "anthropic:claude-haiku-4-5-latest" +CLAUDE_HAIKU_4_5_BEDROCK_MODEL_ID = ( + "bedrock:{region}.anthropic.claude-haiku-4-5-20251001-v1:0" +) +NOVA_LITE_BEDROCK_MODEL_ID = "bedrock:{region}.amazon.nova-2-lite-v1:0" MISTRAL_LARGE_MODEL_ID = "mistral:mistral-large-latest" # ARXIV URLS diff --git a/src/askademic/general.py b/src/askademic/general.py index 042a16b..3ec50ee 100644 --- a/src/askademic/general.py +++ b/src/askademic/general.py @@ -4,7 +4,6 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext -from pydantic_ai.models import Model from pydantic_ai.settings import ModelSettings from askademic.prompts.general import SYSTEM_PROMPT_GENERAL @@ -106,7 +105,7 @@ async def list_research_categories(ctx: RunContext[Context]) -> dict: class GeneralAgent: - def __init__(self, model: Model, model_settings: ModelSettings = None): + def __init__(self, model: str, model_settings: ModelSettings = None): self.agent = general_agent_base self.agent.model = model if model_settings: diff --git a/src/askademic/main.py b/src/askademic/main.py index 899b386..57e5bbb 100644 --- a/src/askademic/main.py +++ b/src/askademic/main.py @@ -48,9 +48,9 @@ async def ask_user_question(): async def check_environment_variables(user_model: str): if user_model == "gemini": - if not os.getenv("GEMINI_API_KEY"): + if not os.getenv("GOOGLE_API_KEY"): console.print( - "[bold red]The GEMINI_API_KEY environment variable is not set.[/bold red]" + "[bold red]The GOOGLE_API_KEY environment variable is not set.[/bold red]" ) sys.exit() elif user_model in "claude": @@ -59,7 +59,7 @@ async def check_environment_variables(user_model: str): "[bold red]The ANTHROPIC_API_KEY environment variable is not set.[/bold red]" ) sys.exit() - elif user_model in ("claude-aws-bedrock", "nova-pro-aws-bedrock"): + elif user_model in ("claude-aws-bedrock", "nova-lite-aws-bedrock"): try: _ = boto3.client("sts").get_caller_identity() except boto3.exceptions.ClientError: @@ -75,7 +75,7 @@ async def check_environment_variables(user_model: str): console.print( "[bold red]Invalid model family selected. " + "Please choose 'gemini', 'claude', 'claude-aws-bedrock'" - + " or 'nova-pro-aws-bedrock'.[/bold red]" + + " or 'nova-lite-aws-bedrock'.[/bold red]" ) sys.exit() @@ -125,12 +125,12 @@ async def ask_me(): "gemini", "claude", "claude-aws-bedrock", - "nova-pro-aws-bedrock", + "nova-lite-aws-bedrock", ): console.print( """[bold red]Please configure the LLM family to be either "gemini" or "claude", "claude-aws-bedrock" - or "nova-pro-aws-bedrock"):[/bold red]""" + or "nova-lite-aws-bedrock"):[/bold red]""" ) return diff --git a/src/askademic/question.py b/src/askademic/question.py index a8eda75..7bf4520 100644 --- a/src/askademic/question.py +++ b/src/askademic/question.py @@ -3,7 +3,6 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent -from pydantic_ai.models import Model from pydantic_ai.settings import ModelSettings from askademic.prompts.general import ( @@ -56,7 +55,7 @@ class QuestionAnswerResponse(BaseModel): class QuestionAgent: def __init__( self, - model: Model, + model: str, model_settings: ModelSettings = None, query_list_limit: int = 10, relevance_score_threshold: float = 0.8, diff --git a/src/askademic/summary.py b/src/askademic/summary.py index 3dd9d18..d8f19bc 100644 --- a/src/askademic/summary.py +++ b/src/askademic/summary.py @@ -3,7 +3,6 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent, Tool -from pydantic_ai.models import Model from pydantic_ai.settings import ModelSettings from askademic.prompts.general import ( @@ -54,14 +53,14 @@ class SummaryResponse(BaseModel): class SummaryAgent: - def __init__(self, model: Model, model_settings: ModelSettings = None): + def __init__(self, model: str, model_settings: ModelSettings = None): self._category_agent = Agent( model=model, model_settings=model_settings, system_prompt=SYSTEM_PROMPT_CATEGORY, output_type=Category, - tools=[Tool(get_categories, takes_ctx=False)], + tools=[Tool(get_categories)], ) self._summary_agent = Agent( @@ -71,7 +70,7 @@ def __init__(self, model: Model, model_settings: ModelSettings = None): output_type=Summary, ) - if "nova" in model.model_name: + if "nova" in str(model): self._max_results = 100 else: self._max_results = 300 diff --git a/src/askademic/utils.py b/src/askademic/utils.py index 1d4115b..37d9687 100644 --- a/src/askademic/utils.py +++ b/src/askademic/utils.py @@ -5,17 +5,13 @@ import boto3 import feedparser import pandas as pd -from pydantic_ai.models import Model -from pydantic_ai.models.anthropic import AnthropicModel -from pydantic_ai.models.bedrock import BedrockConverseModel, BedrockModelSettings -from pydantic_ai.models.gemini import GeminiModel from pydantic_ai.settings import ModelSettings from askademic.constants import ( - CLAUDE_HAIKU_3_5_BEDROCK_MODEL_ID, - CLAUDE_HAIKU_3_5_MODEL_ID, + CLAUDE_HAIKU_4_5_BEDROCK_MODEL_ID, + CLAUDE_HAIKU_4_5_MODEL_ID, GEMINI_2_FLASH_MODEL_ID, - NOVA_PRO_BEDROCK_MODEL_ID, + NOVA_LITE_BEDROCK_MODEL_ID, ) today = datetime.now().strftime("%Y-%m-%d") @@ -23,34 +19,30 @@ logger = logging.getLogger(__name__) -def choose_model(model_family: str = "gemini") -> Tuple[Model, ModelSettings]: +def choose_model(model_family: str = "gemini") -> Tuple[str, ModelSettings]: """ Choose the model ID based on the given model family. + Returns a tuple of (model_string, model_settings) using the provider:model syntax. """ if model_family not in [ "gemini", "claude", "claude-aws-bedrock", - "nova-pro-aws-bedrock", + "nova-lite-aws-bedrock", ]: raise ValueError(f"Invalid model family '{model_family}'.") if model_family == "gemini": - model_name = GEMINI_2_FLASH_MODEL_ID - model = GeminiModel(model_name=model_name) model_settings = ModelSettings(max_tokens=1000, temperature=0) - return model, model_settings + return GEMINI_2_FLASH_MODEL_ID, model_settings elif model_family == "claude": - model_name = CLAUDE_HAIKU_3_5_MODEL_ID - model = AnthropicModel(model_name=model_name) model_settings = ModelSettings(max_tokens=1000, temperature=0) - return model, model_settings - elif model_family in ("claude-aws-bedrock", "nova-pro-aws-bedrock"): - + return CLAUDE_HAIKU_4_5_MODEL_ID, model_settings + elif model_family in ("claude-aws-bedrock", "nova-lite-aws-bedrock"): model_id = ( - CLAUDE_HAIKU_3_5_BEDROCK_MODEL_ID + CLAUDE_HAIKU_4_5_BEDROCK_MODEL_ID if model_family == "claude-aws-bedrock" - else NOVA_PRO_BEDROCK_MODEL_ID + else NOVA_LITE_BEDROCK_MODEL_ID ) region = boto3.session.Session().region_name @@ -60,14 +52,11 @@ def choose_model(model_family: str = "gemini") -> Tuple[Model, ModelSettings]: region = region.split("-")[0] model_name = model_id.format(region=region) - model_settings = BedrockModelSettings( + model_settings = ModelSettings( temperature=0, max_tokens=4000, - top_k=1, - bedrock_performance_configuration={"latency": "optimized"}, ) - model = BedrockConverseModel(model_name=model_name) - return model, model_settings + return model_name, model_settings def list_categories() -> dict: diff --git a/tests/test_article.py b/tests/test_article.py index d724a50..13fd574 100644 --- a/tests/test_article.py +++ b/tests/test_article.py @@ -5,7 +5,7 @@ import pytest from pydantic_ai.agent import AgentRunResult # noqa: F401 -os.environ["GEMINI_API_KEY"] = "mock" +os.environ["GOOGLE_API_KEY"] = "mock" from askademic.article import ( # noqa: E402 USER_PROMPT_ARTICLE_RETRIEVAL_TEMPLATE, @@ -65,7 +65,7 @@ async def test_article_agent( # Mock the agents article_agent = ArticleAgent( - model="gemini-2.0-flash", + model="google-gla:gemini-2.0-flash", ) article_agent._get_article = MagicMock( return_value="Mocked article text for testing purposes." diff --git a/tests/test_question.py b/tests/test_question.py index fd59d8b..353be66 100644 --- a/tests/test_question.py +++ b/tests/test_question.py @@ -5,7 +5,7 @@ import pytest from pydantic_ai.agent import AgentRunResult # noqa: F401 -os.environ["GEMINI_API_KEY"] = "mock" +os.environ["GOOGLE_API_KEY"] = "mock" from askademic.question import ( # noqa: E402 Article, @@ -93,9 +93,9 @@ async def test_question_agent( ): """Test the QuestionAgent class.""" - assert os.environ["GEMINI_API_KEY"] == "mock" + assert os.environ["GOOGLE_API_KEY"] == "mock" - question_agent = QuestionAgent("gemini-2.0-flash") + question_agent = QuestionAgent("google-gla:gemini-2.0-flash") question_agent._query_agent = MagicMock() question_agent._abstract_relevance_agent = MagicMock() question_agent._many_articles_agent = MagicMock() diff --git a/tests/test_summarizer.py b/tests/test_summarizer.py index 5f1a689..4eedf09 100644 --- a/tests/test_summarizer.py +++ b/tests/test_summarizer.py @@ -4,9 +4,8 @@ import pytest from pydantic_ai.agent import AgentRunResult -from pydantic_ai.models.gemini import GeminiModel -os.environ["GEMINI_API_KEY"] = "mock" +os.environ["GOOGLE_API_KEY"] = "mock" from askademic.summary import ( # noqa: E402 Category, @@ -57,7 +56,7 @@ async def test_summary_agent( summary_response, ): """Test the SummaryAgent class.""" - model = GeminiModel(model_name="gemini-2.0-flash") + model = "google-gla:gemini-2.0-flash" summary_agent = SummaryAgent(model) summary_agent._category_agent = MagicMock() summary_agent._summary_agent = MagicMock()