langroid 0.1.85__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/agent/special/relevance_extractor_agent.py

@@ -0,0 +1,127 @@
+"""
+Agent to retrieve relevant segments from a body of text,
+that are relevant to a query.
+
+"""
+
+import logging
+from typing import Optional, no_type_check
+
+from rich.console import Console
+
+from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
+from langroid.agent.chat_document import ChatDocument
+from langroid.agent.tools.segment_extract_tool import SegmentExtractTool
+from langroid.language_models.base import LLMConfig
+from langroid.language_models.openai_gpt import OpenAIGPTConfig
+from langroid.mytypes import Entity
+from langroid.parsing.utils import extract_numbered_segments, number_segments
+from langroid.utils.constants import DONE, NO_ANSWER
+
+console = Console()
+logger = logging.getLogger(__name__)
+
+
+class RelevanceExtractorAgentConfig(ChatAgentConfig):
+    llm: LLMConfig | None = OpenAIGPTConfig()
+    segment_length: int = 1  # number of sentences per segment
+    query: str = ""  # query for relevance extraction
+    system_message = """
+    The user will give you a PASSAGE containing segments numbered as  
+    <#1#>, <#2#>, <#3#>, etc.,
+    followed by a QUERY. Extract ONLY the segment-numbers from 
+    the PASSAGE that are RELEVANT to the QUERY.
+    Present the extracted segment-numbers using the `extract_segments` tool/function.
+    """
+
+
+class RelevanceExtractorAgent(ChatAgent):
+    """
+    Agent for extracting segments from text, that are relevant to a given query.
+    """
+
+    def __init__(self, config: RelevanceExtractorAgentConfig):
+        super().__init__(config)
+        self.config: RelevanceExtractorAgentConfig = config
+        self.enable_message(SegmentExtractTool)
+        self.numbered_passage: Optional[str] = None
+
+    @no_type_check
+    def llm_response(
+        self, message: Optional[str | ChatDocument] = None
+    ) -> Optional[ChatDocument]:
+        """Compose a prompt asking to extract relevant segments from a passage.
+        Steps:
+        - number the segments in the passage
+        - compose prompt
+        - send to LLM
+        """
+        assert self.config.query is not None, "No query specified"
+        assert message is not None, "No message specified"
+        message_str = message.content if isinstance(message, ChatDocument) else message
+        # number the segments in the passage
+        self.numbered_passage = number_segments(message_str, self.config.segment_length)
+        # compose prompt
+        prompt = f"""
+        PASSAGE:
+        {self.numbered_passage}
+        
+        QUERY: {self.config.query}
+        """
+        # send to LLM
+        return super().llm_response(prompt)
+
+    @no_type_check
+    async def llm_response_async(
+        self, message: Optional[str | ChatDocument] = None
+    ) -> Optional[ChatDocument]:
+        """
+        Compose a prompt asking to extract relevant segments from a passage.
+        Steps:
+        - number the segments in the passage
+        - compose prompt
+        - send to LLM
+        The LLM is expected to generate a structured msg according to the
+        SegmentExtractTool schema, i.e. it should contain a `segment_list` field
+        whose value is a list of segment numbers or ranges, like "10,12,14-17".
+        """
+
+        assert self.config.query is not None, "No query specified"
+        assert message is not None, "No message specified"
+        message_str = message.content if isinstance(message, ChatDocument) else message
+        # number the segments in the passage
+        self.numbered_passage = number_segments(message_str, self.config.segment_length)
+        # compose prompt
+        prompt = f"""
+        PASSAGE:
+        {self.numbered_passage}
+        
+        QUERY: {self.config.query}
+        """
+        # send to LLM
+        return await super().llm_response_async(prompt)
+
+    def extract_segments(self, msg: SegmentExtractTool) -> str:
+        """Method to handle a segmentExtractTool message from LLM"""
+        spec = msg.segment_list
+        if len(self.message_history) == 0:
+            return DONE + " " + NO_ANSWER
+        if spec is None or spec.strip() in ["", NO_ANSWER]:
+            return DONE + " " + NO_ANSWER
+        assert self.numbered_passage is not None, "No numbered passage"
+        # assume this has numbered segments
+        try:
+            extracts = extract_numbered_segments(self.numbered_passage, spec)
+        except Exception:
+            return DONE + " " + NO_ANSWER
+        # this response ends the task by saying DONE
+        return DONE + " " + extracts
+
+    def handle_message_fallback(
+        self, msg: str | ChatDocument
+    ) -> str | ChatDocument | None:
+        """Handle case where LLM forgets to use SegmentExtractTool"""
+        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
+            return DONE + " " + NO_ANSWER
+        else:
+            return None

langroid/agent/special/retriever_agent.py

@@ -1,222 +1,56 @@
 """
-Agent to retrieve relevant verbatim whole docs/records from a vector store.
-See test_retriever_agent.py for example usage:
+Deprecated: use DocChatAgent instead, with DocChatAgentConfig.retrieve_only=True,
+and if you want to retrieve FULL relevant doc-contents rather than just extracts,
+then set DocChatAgentConfig.extraction_granularity=-1
+
+This is an agent to retrieve relevant extracts from a vector store,
+where the LLM is used to filter for "true" relevance after retrieval from the
+vector store.
+This is essentially the same as DocChatAgent, except that instead of
+generating final summary answer based on relevant extracts, it just returns
+those extracts.
+See test_retriever_agent.py for example usage.
 """
+
 import logging
-from abc import ABC, abstractmethod
-from typing import List, Optional, Sequence
+from typing import Sequence
 
-from rich import print
 from rich.console import Console
 
-from langroid.agent.chat_document import ChatDocMetaData, ChatDocument
 from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
-from langroid.embedding_models.models import OpenAIEmbeddingsConfig
-from langroid.language_models.base import StreamingIfAllowed
-from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
-from langroid.mytypes import DocMetaData, Document, Entity
-from langroid.parsing.parser import ParsingConfig, Splitter
-from langroid.prompts.prompts_config import PromptsConfig
-from langroid.utils.constants import NO_ANSWER
-from langroid.vector_store.base import VectorStoreConfig
-from langroid.vector_store.qdrantdb import QdrantDBConfig
+from langroid.mytypes import DocMetaData, Document
 
 console = Console()
 logger = logging.getLogger(__name__)
 
+# for backwards compatibility:
+RecordMetadata = DocMetaData
+RecordDoc = Document
+RetrieverAgentConfig = DocChatAgentConfig
 
-class RecordMetadata(DocMetaData):
-    id: None | int | str = None
-
-
-class RecordDoc(Document):
-    metadata: RecordMetadata
-
-
-class RetrieverAgentConfig(DocChatAgentConfig):
-    n_matches: int = 3
-    debug: bool = False
-    max_context_tokens = 500
-    conversation_mode = True
-    cache: bool = True  # cache results
-    gpt4: bool = True  # use GPT-4
-    stream: bool = True  # allow streaming where needed
-    max_tokens: int = 10000
-    vecdb: VectorStoreConfig = QdrantDBConfig(
-        collection_name=None,
-        storage_path=".qdrant/data/",
-        embedding=OpenAIEmbeddingsConfig(
-            model_type="openai",
-            model_name="text-embedding-ada-002",
-            dims=1536,
-        ),
-    )
 
-    llm: OpenAIGPTConfig = OpenAIGPTConfig(
-        type="openai",
-        chat_model=OpenAIChatModel.GPT4,
-    )
-    parsing: ParsingConfig = ParsingConfig(
-        splitter=Splitter.TOKENS,
-        chunk_size=100,
-        n_similar_docs=5,
-    )
-
-    prompts: PromptsConfig = PromptsConfig(
-        max_tokens=1000,
-    )
-
-
-class RetrieverAgent(DocChatAgent, ABC):
+class RetrieverAgent(DocChatAgent):
     """
-    Agent for retrieving whole records/docs matching a query
+    Agent for just retrieving chunks/docs/extracts matching a query
     """
 
-    def __init__(self, config: RetrieverAgentConfig):
+    def __init__(self, config: DocChatAgentConfig):
         super().__init__(config)
-        self.config: RetrieverAgentConfig = config
+        self.config: DocChatAgentConfig = config
+        logger.warning(
+            """
+        `RetrieverAgent` is deprecated. Use `DocChatAgent` instead, with
+        `DocChatAgentConfig.retrieve_only=True`, and if you want to retrieve
+        FULL relevant doc-contents rather than just extracts, then set
+        `DocChatAgentConfig.extraction_granularity=-1`
+        """
+        )
 
-    @abstractmethod
-    def get_records(self) -> Sequence[RecordDoc]:
-        pass
+    def get_records(self) -> Sequence[Document]:
+        raise NotImplementedError
 
     def ingest(self) -> None:
         records = self.get_records()
         if self.vecdb is None:
             raise ValueError("No vector store specified")
         self.vecdb.add_documents(records)
-
-    def llm_response(
-        self,
-        query: None | str | ChatDocument = None,
-    ) -> Optional[ChatDocument]:
-        if not self.llm_can_respond(query):
-            return None
-        if query is None:
-            return super().llm_response(None)  # type: ignore
-        if isinstance(query, ChatDocument):
-            query_str = query.content
-        else:
-            query_str = query
-        docs = self.get_relevant_extracts(query_str)
-        if len(docs) == 0:
-            return None
-        content = "\n\n".join([d.content for d in docs])
-        print(f"[green]{content}")
-        meta = dict(
-            sender=Entity.LLM,
-        )
-        meta.update(docs[0].metadata)
-
-        return ChatDocument(
-            content=content,
-            metadata=ChatDocMetaData(**meta),
-        )
-
-    def get_nearest_docs(self, query: str) -> List[Document]:
-        """
-        Given a query, get the records/docs whose contents are closest to the
-            query, in terms of vector similarity.
-        Args:
-            query: query string
-        Returns:
-            list of Document objects
-        """
-        if self.vecdb is None:
-            logger.warning("No vector store specified")
-            return []
-        with console.status("[cyan]Searching VecDB for similar docs/records..."):
-            docs_and_scores = self.vecdb.similar_texts_with_scores(
-                query,
-                k=self.config.parsing.n_similar_docs,
-            )
-        docs: List[Document] = [
-            Document(content=d.content, metadata=d.metadata)
-            for (d, _) in docs_and_scores
-        ]
-        return docs
-
-    def get_relevant_extracts(self, query: str) -> List[Document]:
-        """
-        Given a query, get the records/docs whose contents are most relevant to the
-            query. First get nearest docs from vector store, then select the best
-            matches according to the LLM.
-        Args:
-            query (str): query string
-
-        Returns:
-            List[Document]: list of Document objects
-        """
-        response = Document(
-            content=NO_ANSWER,
-            metadata=DocMetaData(
-                source="None",
-            ),
-        )
-        nearest_docs = self.get_nearest_docs(query)
-        if len(nearest_docs) == 0:
-            return [response]
-        if self.llm is None:
-            logger.warning("No LLM specified")
-            return nearest_docs
-        with console.status("LLM selecting relevant docs from retrieved ones..."):
-            with StreamingIfAllowed(self.llm, False):
-                doc_list = self.llm_select_relevant_docs(query, nearest_docs)
-
-        return doc_list
-
-    def llm_select_relevant_docs(
-        self, query: str, docs: List[Document]
-    ) -> List[Document]:
-        """
-        Given a query and a list of docs, select the docs whose contents match best,
-            according to the LLM. Use the doc IDs to select the docs from the vector
-            store.
-        Args:
-            query: query string
-            docs: list of Document objects
-        Returns:
-            list of Document objects
-        """
-        doc_contents = "\n\n".join(
-            [f"DOC: ID={d.id()}, CONTENT: {d.content}" for d in docs]
-        )
-        prompt = f"""
-        Given the following QUERY: 
-        {query}
-        and the following DOCS with IDs and contents
-        {doc_contents}
-        
-        Find at most {self.config.n_matches} DOCs that are most relevant to the QUERY.
-        Return your answer as a sequence of DOC IDS ONLY, for example: 
-        "id1 id2 id3..."
-        If there are no relevant docs, simply say {NO_ANSWER}.
-        Even if there is only one relevant doc, return it as a single ID.
-        Do not give any explanations or justifications.
-        """
-        default_response = Document(
-            content=NO_ANSWER,
-            metadata=DocMetaData(
-                source="None",
-            ),
-        )
-
-        if self.llm is None:
-            logger.warning("No LLM specified")
-            return [default_response]
-        response = self.llm.generate(
-            prompt, max_tokens=self.config.llm.max_output_tokens
-        )
-        if response.message == NO_ANSWER:
-            return [default_response]
-        ids = response.message.split()
-        if len(ids) == 0:
-            return [default_response]
-        if self.vecdb is None:
-            logger.warning("No vector store specified")
-            return [default_response]
-        docs = self.vecdb.get_documents_by_ids(ids)
-        return [
-            Document(content=d.content, metadata=DocMetaData(source="LLM"))
-            for d in docs
-        ]

langroid/agent/special/sql/__init__.py

@@ -0,0 +1,11 @@
+from .sql_chat_agent import SQLChatAgentConfig, SQLChatAgent
+
+from . import sql_chat_agent
+from . import utils
+
+__all__ = [
+    "SQLChatAgentConfig",
+    "SQLChatAgent",
+    "sql_chat_agent",
+    "utils",
+]

langroid/agent/special/sql/sql_chat_agent.py

@@ -6,12 +6,13 @@ Functionality includes:
 - adding table and column context
 - asking a question about a SQL schema
 """
+
 import logging
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 from rich import print
 from rich.console import Console
-from sqlalchemy import MetaData, Row, create_engine, text
+from sqlalchemy import MetaData, Row, create_engine, inspect, text
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import Session, sessionmaker
@@ -35,9 +36,7 @@ from langroid.agent.special.sql.utils.tools import (
     GetTableSchemaTool,
     RunQueryTool,
 )
-from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
 from langroid.mytypes import Entity
-from langroid.prompts.prompts_config import PromptsConfig
 from langroid.vector_store.base import VectorStoreConfig
 
 logger = logging.getLogger(__name__)
@@ -67,7 +66,6 @@ SQL_ERROR_MSG = "There was an error in your SQL Query"
 class SQLChatAgentConfig(ChatAgentConfig):
     system_message: str = DEFAULT_SQL_CHAT_SYSTEM_MESSAGE
     user_message: None | str = None
-    max_context_tokens: int = 1000
     cache: bool = True  # cache results
     debug: bool = False
     stream: bool = True  # allow streaming where needed
@@ -76,6 +74,7 @@ class SQLChatAgentConfig(ChatAgentConfig):
     vecdb: None | VectorStoreConfig = None
     context_descriptions: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = {}
     use_schema_tools: bool = False
+    multi_schema: bool = False
 
     """
     Optional, but strongly recommended, context descriptions for tables, columns, 
@@ -90,6 +89,9 @@ class SQLChatAgentConfig(ChatAgentConfig):
     is another table name and the value is a description of the relationship to 
     that table.
 
+    If multi_schema support is enabled, the tables names in the description
+    should be of the form 'schema_name.table_name'.
+
     For example:
     {
         'table1': {
@@ -109,15 +111,6 @@ class SQLChatAgentConfig(ChatAgentConfig):
     }
     """
 
-    llm: OpenAIGPTConfig = OpenAIGPTConfig(
-        type="openai",
-        chat_model=OpenAIChatModel.GPT4,
-        completion_model=OpenAIChatModel.GPT4,
-    )
-    prompts: PromptsConfig = PromptsConfig(
-        max_tokens=1000,
-    )
-
 
 class SQLChatAgent(ChatAgent):
     """
@@ -155,19 +148,44 @@ class SQLChatAgent(ChatAgent):
         """Initialize the database metadata."""
         if self.engine is None:
             raise ValueError("Database engine is None")
+        self.metadata: MetaData | List[MetaData] = []
 
-        self.metadata = MetaData()
-        self.metadata.reflect(self.engine)
-        logger.info(
-            "SQLChatAgent initialized with database: %s and tables: %s",
-            self.engine,
-            self.metadata.tables,
-        )
+        if self.config.multi_schema:
+            logger.info(
+                "Initializing SQLChatAgent with database: %s",
+                self.engine,
+            )
+
+            self.metadata = []
+            inspector = inspect(self.engine)
+
+            for schema in inspector.get_schema_names():
+                metadata = MetaData(schema=schema)
+                metadata.reflect(self.engine)
+                self.metadata.append(metadata)
+
+                logger.info(
+                    "Initializing SQLChatAgent with database: %s, schema: %s, "
+                    "and tables: %s",
+                    self.engine,
+                    schema,
+                    metadata.tables,
+                )
+        else:
+            self.metadata = MetaData()
+            self.metadata.reflect(self.engine)
+            logger.info(
+                "SQLChatAgent initialized with database: %s and tables: %s",
+                self.engine,
+                self.metadata.tables,
+            )
 
     def _init_table_metadata(self) -> None:
         """Initialize metadata for the tables present in the database."""
         if not self.config.context_descriptions and isinstance(self.engine, Engine):
-            self.config.context_descriptions = extract_schema_descriptions(self.engine)
+            self.config.context_descriptions = extract_schema_descriptions(
+                self.engine, self.config.multi_schema
+            )
 
         if self.config.use_schema_tools:
             self.table_metadata = populate_metadata_with_schema_tools(
@@ -228,8 +246,10 @@ class SQLChatAgent(ChatAgent):
         if isinstance(msg, ChatDocument) and msg.function_call is not None:
             sender_name = msg.function_call.name
 
+        content = results.content if isinstance(results, ChatDocument) else results
+
         return ChatDocument(
-            content=results,
+            content=content,
             metadata=ChatDocMetaData(
                 source=Entity.AGENT,
                 sender=Entity.AGENT,
@@ -329,6 +349,10 @@ class SQLChatAgent(ChatAgent):
         Returns:
             str: The names of all tables in the database.
         """
+        if isinstance(self.metadata, list):
+            table_names = [", ".join(md.tables.keys()) for md in self.metadata]
+            return ", ".join(table_names)
+
         return ", ".join(self.metadata.tables.keys())
 
     def get_table_schema(self, msg: GetTableSchemaTool) -> str:

langroid/agent/special/sql/utils/__init__.py

@@ -0,0 +1,22 @@
+from .tools import (
+    RunQueryTool,
+    GetTableNamesTool,
+    GetTableSchemaTool,
+    GetColumnDescriptionsTool,
+)
+
+from . import description_extractors
+from . import populate_metadata
+from . import system_message
+from . import tools
+
+__all__ = [
+    "RunQueryTool",
+    "GetTableNamesTool",
+    "GetTableSchemaTool",
+    "GetColumnDescriptionsTool",
+    "description_extractors",
+    "populate_metadata",
+    "system_message",
+    "tools",
+]