149 rehaul update pydantic ai

.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"

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
 

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

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
 

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",
+]

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

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

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:

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
 

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,

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

src/askademic/utils.py

@@ -5,52 +5,44 @@
 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")
 logging.basicConfig(level=logging.INFO, filename=f"logs/{today}_logs.txt")
 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:

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."

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()

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()