Propagates llmclient changes to pqa (#839)

Author: maykcaldas

.mailmap

@@ -1,12 +1,13 @@
 Andrew White <[email protected]> <[email protected]>
-James Braza <[email protected]> <jamesbraza@gmail.com>
+Anush008 <[email protected]> Anush <anushshetty90@gmail.com>
 Geemi Wellawatte <[email protected]> <[email protected]>
 Geemi Wellawatte <[email protected]> <[email protected]>
+Harry Vu <[email protected]> <[email protected]>
+Harry Vu <[email protected]> harryvu-futurehouse
+James Braza <[email protected]> <[email protected]>
+Mayk Caldas <[email protected]> maykcaldas
+Mayk Caldas <[email protected]> <[email protected]>
 Michael Skarlinski <[email protected]> mskarlin <[email protected]>
 Odhran O'Donoghue <[email protected]> odhran-o-d <[email protected]>
 Odhran O'Donoghue <[email protected]> <[email protected]>
 Samantha Cox <[email protected]> <[email protected]>
-Anush008 <[email protected]> Anush <[email protected]>
-Mayk Caldas <[email protected]> maykcaldas <[email protected]>
-Harry Vu <[email protected]> <[email protected]>
-Harry Vu <[email protected]> harryvu-futurehouse

.pre-commit-config.yaml

@@ -84,9 +84,9 @@ repos:
           - aiohttp>=3.10.6 # Match pyproject.toml
           - PyMuPDF>=1.24.12
           - anyio
-          - fh-llm-client[deepseek]<0.1.0 # Match pyproject.toml
+          - fhlmi>=0.0.1
           - fhaviary[llm]>=0.18.2 # Match pyproject.toml
-          - ldp>=0.20 # Match pyproject.toml
+          - ldp>=0.25.0 # Match pyproject.toml
           - html2text
           - httpx
           - pybtex

paperqa/__init__.py

@@ -1,6 +1,6 @@
 import warnings
 
-from llmclient import (
+from lmi import (
     EmbeddingModel,
     HybridEmbeddingModel,
     LiteLLMEmbeddingModel,

paperqa/agents/env.py

@@ -13,7 +13,7 @@
     ToolResponseMessage,
 )
 from aviary.utils import MultipleChoiceQuestion
-from llmclient import EmbeddingModel, LiteLLMModel
+from lmi import EmbeddingModel, LiteLLMModel
 
 from paperqa.docs import Docs
 from paperqa.settings import Settings

paperqa/agents/helpers.py

@@ -5,7 +5,8 @@
 from datetime import datetime
 from typing import cast
 
-from llmclient import LiteLLMModel, LLMModel
+from aviary.core import Message
+from lmi import LiteLLMModel, LLMModel
 from rich.table import Table
 
 from paperqa.docs import Docs
@@ -60,12 +61,13 @@ async def litellm_get_search_query(
         )
     else:
         model = llm
-    result = await model.run_prompt(
-        prompt=search_prompt,
-        data={"question": question, "count": count},
-        system_prompt=None,
+    messages = [
+        Message(content=search_prompt.format(question=question, count=count)),
+    ]
+    result = await model.call_single(
+        messages=messages,
     )
-    search_query = result.text
+    search_query = cast(str, result.text)
     queries = [s for s in search_query.split("\n") if len(s) > 3]  # noqa: PLR2004
     # remove "2.", "3.", etc. -- https://regex101.com/r/W2f7F1/1
     queries = [re.sub(r"^\d+\.\s*", "", q) for q in queries]

paperqa/agents/models.py

@@ -5,10 +5,11 @@
 import time
 from contextlib import asynccontextmanager
 from enum import StrEnum
-from typing import Any, ClassVar, Protocol
+from typing import Any, ClassVar, Protocol, cast
 from uuid import UUID, uuid4
 
-from llmclient import LiteLLMModel, LLMModel
+from aviary.core import Message
+from lmi import LiteLLMModel, LLMModel
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -79,12 +80,20 @@ async def get_summary(self, llm_model: LLMModel | str = "gpt-4o") -> str:
         model = (
             LiteLLMModel(name=llm_model) if isinstance(llm_model, str) else llm_model
         )
-        result = await model.run_prompt(
-            prompt="{question}\n\n{answer}",
-            data={"question": self.session.question, "answer": self.session.answer},
-            system_prompt=sys_prompt,
+        prompt_template = "{question}\n\n{answer}"
+        messages = [
+            Message(role="system", content=sys_prompt),
+            Message(
+                role="user",
+                content=prompt_template.format(
+                    question=self.session.question, answer=self.session.answer
+                ),
+            ),
+        ]
+        result = await model.call_single(
+            messages=messages,
         )
-        return result.text.strip()
+        return cast(str, result.text).strip()
 
 
 class TimerData(BaseModel):

paperqa/agents/tools.py

@@ -11,7 +11,7 @@
 from typing import ClassVar, Self, cast
 
 from aviary.core import ToolRequestMessage
-from llmclient import Embeddable, EmbeddingModel, LiteLLMModel
+from lmi import Embeddable, EmbeddingModel, LiteLLMModel
 from pydantic import BaseModel, ConfigDict, Field, computed_field
 
 from paperqa.docs import Docs

paperqa/clients/__init__.py

@@ -6,10 +6,10 @@
 from typing import Any, cast
 
 import aiohttp
+from lmi.utils import gather_with_concurrency
 from pydantic import BaseModel, ConfigDict
 
 from paperqa.types import Doc, DocDetails
-from paperqa.utils import gather_with_concurrency
 
 from .client_models import MetadataPostProcessor, MetadataProvider
 from .crossref import CrossrefProvider

paperqa/core.py

@@ -3,9 +3,11 @@
 import json
 import re
 from collections.abc import Callable, Sequence
-from typing import Any
+from typing import Any, cast
+
+from aviary.core import Message
+from lmi import LLMModel
 
-from paperqa.llms import PromptRunner
 from paperqa.types import Context, LLMResult, Text
 from paperqa.utils import extract_score, strip_citations
 
@@ -102,7 +104,8 @@ def fraction_replacer(match: re.Match) -> str:
 async def map_fxn_summary(
     text: Text,
     question: str,
-    prompt_runner: PromptRunner | None,
+    summary_llm_model: LLMModel | None,
+    prompt_templates: tuple[str, str] | None,
     extra_prompt_data: dict[str, str] | None = None,
     parser: Callable[[str], dict[str, Any]] | None = None,
     callbacks: Sequence[Callable[[str], None]] | None = None,
@@ -115,12 +118,14 @@ async def map_fxn_summary(
 
     Args:
         text: The text to parse.
-        question: The question to use for the chain.
-        prompt_runner: The prompt runner to call - should have question, citation,
-            summary_length, and text fields.
-        extra_prompt_data: Optional extra kwargs to pass to the prompt runner's data.
-        parser: The parser to use for parsing - return empty dict on Failure to fallback to text parsing.
-        callbacks: LLM callbacks to execute in the prompt runner.
+        question: The question to use for summarization.
+        summary_llm_model: The LLM model to use for generating summaries.
+        prompt_templates: Optional two-elements tuple containing templates for the user and system prompts.
+            prompt_templates = (user_prompt_template, system_prompt_template)
+        extra_prompt_data: Optional extra data to pass to the prompt template.
+        parser: Optional parser function to parse LLM output into structured data.
+            Should return dict with at least 'summary' field.
+        callbacks: Optional sequence of callback functions to execute during LLM calls.
 
     Returns:
         The context object and LLMResult to get info about the LLM execution.
@@ -131,14 +136,21 @@ async def map_fxn_summary(
     citation = text.name + ": " + text.doc.formatted_citation
     success = False
 
-    if prompt_runner:
-        llm_result = await prompt_runner(
-            {"question": question, "citation": citation, "text": text.text}
-            | (extra_prompt_data or {}),
-            callbacks,
-            "evidence:" + text.name,
+    if summary_llm_model and prompt_templates:
+        data = {"question": question, "citation": citation, "text": text.text} | (
+            extra_prompt_data or {}
+        )
+        message_prompt, system_prompt = prompt_templates
+        messages = [
+            Message(role="system", content=system_prompt.format(**data)),
+            Message(role="user", content=message_prompt.format(**data)),
+        ]
+        llm_result = await summary_llm_model.call_single(
+            messages=messages,
+            callbacks=callbacks,
+            name="evidence:" + text.name,
         )
-        context = llm_result.text
+        context = cast(str, llm_result.text)
         result_data = parser(context) if parser else {}
         success = bool(result_data)
         if success: