langroid 0.58.2__py3-none-any.whl → 0.59.0b1__py3-none-any.whl

langroid/agent/special/arangodb/tools.py

@@ -13,8 +13,8 @@ class AQLRetrievalTool(ToolMessage):
     """
     aql_query: str
 
-    _max_result_tokens = 500
-    _max_retained_tokens = 200
+    _max_result_tokens: int = 500
+    _max_retained_tokens: int = 200
 
     @classmethod
     def examples(cls) -> List[ToolMessage | Tuple[str, ToolMessage]]:
@@ -101,7 +101,7 @@ class ArangoSchemaTool(ToolMessage):
     properties: bool = True
     collections: List[str] | None = None
 
-    _max_result_tokens = 500
+    _max_result_tokens: int = 500
 
 
 arango_schema_tool_name = ArangoSchemaTool.default_value("request")

langroid/agent/special/doc_chat_agent.py

@@ -388,14 +388,14 @@ class DocChatAgent(ChatAgent):
                     p: (
                         m
                         if isinstance(m, dict)
-                        else (isinstance(m, DocMetaData) and m.dict())
+                        else (isinstance(m, DocMetaData) and m.model_dump())
                     )  # appease mypy
                     for p, m in zip(idxs, metadata)
                 }
             elif isinstance(metadata, dict):
                 idx2meta = {p: metadata for p in idxs}
             else:
-                idx2meta = {p: metadata.dict() for p in idxs}
+                idx2meta = {p: metadata.model_dump() for p in idxs}
             urls_meta = {u: idx2meta[u] for u in url_idxs}
             paths_meta = {p: idx2meta[p] for p in path_idxs}
         docs: List[Document] = []
@@ -412,7 +412,7 @@ class DocChatAgent(ChatAgent):
                 # update metadata of each doc with meta
                 for d in url_docs:
                     orig_source = d.metadata.source
-                    d.metadata = d.metadata.copy(update=meta)
+                    d.metadata = d.metadata.model_copy(update=meta)
                     d.metadata.source = _append_metadata_source(
                         orig_source, meta.get("source", "")
                     )
@@ -429,7 +429,7 @@ class DocChatAgent(ChatAgent):
                 # update metadata of each doc with meta
                 for d in path_docs:
                     orig_source = d.metadata.source
-                    d.metadata = d.metadata.copy(update=meta)
+                    d.metadata = d.metadata.model_copy(update=meta)
                     d.metadata.source = _append_metadata_source(
                         orig_source, meta.get("source", "")
                     )
@@ -474,22 +474,22 @@ class DocChatAgent(ChatAgent):
         if isinstance(metadata, list) and len(metadata) > 0:
             for d, m in zip(docs, metadata):
                 orig_source = d.metadata.source
-                m_dict = m if isinstance(m, dict) else m.dict()  # type: ignore
-                d.metadata = d.metadata.copy(update=m_dict)  # type: ignore
+                m_dict = m if isinstance(m, dict) else m.model_dump()  # type: ignore
+                d.metadata = d.metadata.model_copy(update=m_dict)  # type: ignore
                 d.metadata.source = _append_metadata_source(
                     orig_source, m_dict.get("source", "")
                 )
         elif isinstance(metadata, dict):
             for d in docs:
                 orig_source = d.metadata.source
-                d.metadata = d.metadata.copy(update=metadata)
+                d.metadata = d.metadata.model_copy(update=metadata)
                 d.metadata.source = _append_metadata_source(
                     orig_source, metadata.get("source", "")
                 )
         elif isinstance(metadata, DocMetaData):
             for d in docs:
                 orig_source = d.metadata.source
-                d.metadata = d.metadata.copy(update=metadata.dict())
+                d.metadata = d.metadata.model_copy(update=metadata.model_dump())
                 d.metadata.source = _append_metadata_source(
                     orig_source, metadata.source
                 )
@@ -1025,10 +1025,12 @@ class DocChatAgent(ChatAgent):
                     f"{doc.content}{enrichment_config.delimiter}{enrichment}"
                 )
 
-                new_doc = doc.copy(
+                new_doc = doc.model_copy(
                     update={
                         "content": combined_content,
-                        "metadata": doc.metadata.copy(update={"has_enrichment": True}),
+                        "metadata": doc.metadata.model_copy(
+                            update={"has_enrichment": True}
+                        ),
                     }
                 )
                 augmented_docs.append(new_doc)
@@ -1212,7 +1214,7 @@ class DocChatAgent(ChatAgent):
             return docs_scores
         if len(docs_scores) == 0:
             return []
-        if set(docs_scores[0][0].__fields__) != {"content", "metadata"}:
+        if set(docs_scores[0][0].model_fields) != {"content", "metadata"}:
             # Do not add context window when there are other fields besides just
             # content and metadata, since we do not know how to set those other fields
             # for newly created docs with combined content.
@@ -1515,7 +1517,7 @@ class DocChatAgent(ChatAgent):
         delimiter = self.config.chunk_enrichment_config.delimiter
         return [
             (
-                doc.copy(update={"content": doc.content.split(delimiter)[0]})
+                doc.model_copy(update={"content": doc.content.split(delimiter)[0]})
                 if doc.content and getattr(doc.metadata, "has_enrichment", False)
                 else doc
             )
@@ -1572,7 +1574,7 @@ class DocChatAgent(ChatAgent):
         for p, e in zip(passages, extracts):
             if e == NO_ANSWER or len(e) == 0:
                 continue
-            p_copy = p.copy()
+            p_copy = p.model_copy()
             p_copy.content = e
             passage_extracts.append(p_copy)
 
@@ -1607,7 +1609,7 @@ class DocChatAgent(ChatAgent):
                 sender=Entity.LLM,
             )
             # copy metadata from first doc, unclear what to do here.
-            meta.update(extracts[0].metadata)
+            meta.update(extracts[0].metadata.model_dump())
             return ChatDocument(
                 content="\n\n".join([e.content for e in extracts]),
                 metadata=ChatDocMetaData(**meta),  # type: ignore

langroid/agent/special/lance_rag/critic_agent.py

@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
 
 
 class QueryPlanCriticConfig(LanceQueryPlanAgentConfig):
-    name = "QueryPlanCritic"
-    system_message = f"""
+    name: str = "QueryPlanCritic"
+    system_message: str = f"""
     You are an expert at carefully planning a query that needs to be answered
     based on a large collection of documents. These docs have a special `content` field
     and additional FILTERABLE fields in the SCHEMA below, along with the 

langroid/agent/special/lance_rag/query_planner_agent.py

@@ -37,11 +37,11 @@ class LanceQueryPlanAgentConfig(ChatAgentConfig):
     critic_name: str = "QueryPlanCritic"
     doc_agent_name: str = "LanceRAG"
     doc_schema: str = ""
-    use_tools = False
+    use_tools: bool = False
     max_retries: int = 5  # max number of retries for query plan
-    use_functions_api = True
+    use_functions_api: bool = True
 
-    system_message = """
+    system_message: str = """
     You will receive a QUERY, to be answered based on an EXTREMELY LARGE collection
     of documents you DO NOT have access to, but your ASSISTANT does.
     You only know that these documents have a special `content` field
@@ -174,7 +174,7 @@ class LanceQueryPlanAgent(ChatAgent):
 
         # (a) insert `recipient` in the QueryPlanTool:
         # QPWithRecipient = QueryPlanTool.require_recipient()
-        # qp = QPWithRecipient(**msg.dict(), recipient=self.config.doc_agent_name)
+        # qp = QPWithRecipient(**msg.model_dump(), recipient=self.config.doc_agent_name)
         # return qp
         #
         # OR

langroid/agent/special/lance_tools.py

@@ -1,7 +1,8 @@
 import logging
 
+from pydantic import BaseModel, Field
+
 from langroid.agent.tool_message import ToolMessage
-from langroid.pydantic_v1 import BaseModel, Field
 
 logger = logging.getLogger(__name__)
 
@@ -19,8 +20,8 @@ class QueryPlan(BaseModel):
 
 
 class QueryPlanTool(ToolMessage):
-    request = "query_plan"  # the agent method name that handles this tool
-    purpose = """
+    request: str = "query_plan"  # the agent method name that handles this tool
+    purpose: str = """
     Given a user's query, generate a query <plan> consisting of:
     - <original_query> - the original query for reference
     - <filter> condition if needed (or empty string if no filter is needed)
@@ -52,8 +53,8 @@ class QueryPlanAnswerTool(ToolMessage):
 
 
 class QueryPlanFeedbackTool(ToolMessage):
-    request = "query_plan_feedback"
-    purpose = """
+    request: str = "query_plan_feedback"
+    purpose: str = """
     To give <feedback> regarding the query plan, 
     along with a <suggested_fix> if any (empty string if no fix is suggested).
     """

langroid/agent/special/lance_tools.py-e

@@ -0,0 +1,61 @@
+import logging
+
+from langroid.agent.tool_message import ToolMessage
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+
+class QueryPlan(BaseModel):
+    original_query: str = Field(..., description="The original query for reference")
+    query: str = Field(..., description="A possibly NON-EMPTY rephrased query")
+    filter: str = Field(
+        "",
+        description="Filter condition if needed (or empty if no filter is needed)",
+    )
+    dataframe_calc: str = Field(
+        "", description="An optional Pandas-dataframe calculation/aggregation string"
+    )
+
+
+class QueryPlanTool(ToolMessage):
+    request = "query_plan"  # the agent method name that handles this tool
+    purpose = """
+    Given a user's query, generate a query <plan> consisting of:
+    - <original_query> - the original query for reference
+    - <filter> condition if needed (or empty string if no filter is needed)
+    - <query> - a possibly NON-EMPTY rephrased query that can be used to match the 
+        CONTENT of the documents 
+        (can be same as <original_query> if no rephrasing is needed)
+    - <dataframe_calc> - a Pandas-dataframe calculation/aggregation string
+        that can be used to calculate the answer 
+        (or empty string if no calculation is needed).
+    """
+    plan: QueryPlan
+
+
+class AnswerTool(ToolMessage):
+    """Wrapper for answer from LanceDocChatAgent"""
+
+    purpose: str = "To package the answer from LanceDocChatAgent"
+    request: str = "answer_tool"
+    answer: str
+
+
+class QueryPlanAnswerTool(ToolMessage):
+    request: str = "query_plan_answer"  # the agent method name that handles this tool
+    purpose: str = """
+    Assemble query <plan> and <answer>
+    """
+    plan: QueryPlan
+    answer: str = Field(..., description="The answer received from the assistant")
+
+
+class QueryPlanFeedbackTool(ToolMessage):
+    request = "query_plan_feedback"
+    purpose = """
+    To give <feedback> regarding the query plan, 
+    along with a <suggested_fix> if any (empty string if no fix is suggested).
+    """
+    feedback: str
+    suggested_fix: str

langroid/agent/special/neo4j/neo4j_chat_agent.py

@@ -1,11 +1,11 @@
 import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings, SettingsConfigDict
 from rich import print
 from rich.console import Console
 
-from langroid.pydantic_v1 import BaseModel, BaseSettings
-
 if TYPE_CHECKING:
     import neo4j
 
@@ -47,11 +47,7 @@ class Neo4jSettings(BaseSettings):
     password: str = ""
     database: str = ""
 
-    class Config:
-        # This enables the use of environment variables to set the settings,
-        # e.g. NEO4J_URI, NEO4J_USERNAME, etc.,
-        # which can either be set in a .env file or in the shell via export cmds.
-        env_prefix = "NEO4J_"
+    model_config = SettingsConfigDict(env_prefix="NEO4J_")
 
 
 class QueryResult(BaseModel):