langroid 0.33.6__py3-none-any.whl → 0.33.7__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

@@ -0,0 +1,56 @@
+"""
+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 typing import Sequence
+
+from rich.console import Console
+
+from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
+from langroid.mytypes import DocMetaData, Document
+
+console = Console()
+logger = logging.getLogger(__name__)
+
+# for backwards compatibility:
+RecordMetadata = DocMetaData
+RecordDoc = Document
+RetrieverAgentConfig = DocChatAgentConfig
+
+
+class RetrieverAgent(DocChatAgent):
+    """
+    Agent for just retrieving chunks/docs/extracts matching a query
+    """
+
+    def __init__(self, config: DocChatAgentConfig):
+        super().__init__(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`
+        """
+        )
+
+    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)

langroid/agent/special/sql/__init__.py

@@ -0,0 +1,17 @@
+from . import utils
+
+
+__all__ = [
+    "utils",
+]
+
+try:
+    from . import sql_chat_agent
+    from .sql_chat_agent import SQLChatAgentConfig, SQLChatAgent
+
+    sql_chat_agent
+    SQLChatAgent
+    SQLChatAgentConfig
+    __all__.extend(["SQLChatAgentConfig", "SQLChatAgent", "sql_chat_agent"])
+except ImportError:
+    pass