langroid 0.58.3__py3-none-any.whl → 0.59.0__py3-none-any.whl

langroid/agent/special/arangodb/arangodb_agent.py

@@ -8,6 +8,8 @@ from arango.client import ArangoClient
 from arango.database import StandardDatabase
 from arango.exceptions import ArangoError, ServerConnectionError
 from numpy import ceil
+from pydantic import BaseModel, ConfigDict
+from pydantic_settings import BaseSettings, SettingsConfigDict
 from rich import print
 from rich.console import Console
 
@@ -31,7 +33,6 @@ from langroid.agent.special.arangodb.utils import count_fields, trim_schema
 from langroid.agent.tools.orchestration import DoneTool, ForwardTool
 from langroid.exceptions import LangroidImportError
 from langroid.mytypes import Entity
-from langroid.pydantic_v1 import BaseModel, BaseSettings
 from langroid.utils.constants import SEND_TO
 
 logger = logging.getLogger(__name__)
@@ -49,8 +50,7 @@ class ArangoSettings(BaseSettings):
     password: str = ""
     database: str = ""
 
-    class Config:
-        env_prefix = "ARANGO_"
+    model_config = SettingsConfigDict(env_prefix="ARANGO_")
 
 
 class QueryResult(BaseModel):
@@ -68,22 +68,14 @@ class QueryResult(BaseModel):
         ]
     ] = None
 
-    class Config:
-        # Allow arbitrary types for flexibility
-        arbitrary_types_allowed = True
-
-        # Handle JSON serialization of special types
-        json_encoders = {
-            # Add custom encoders if needed, e.g.:
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        json_encoders={
             datetime.datetime: lambda v: v.isoformat(),
-            # Could add others for specific ArangoDB types
-        }
-
-        # Validate all assignments
-        validate_assignment = True
-
-        # Frozen=True if we want immutability
-        frozen = False
+        },
+        validate_assignment=True,
+        frozen=False,
+    )
 
 
 class ArangoChatAgentConfig(ChatAgentConfig):

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/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):

langroid/agent/special/relevance_extractor_agent.py

@@ -26,7 +26,7 @@ class RelevanceExtractorAgentConfig(ChatAgentConfig):
     llm: LLMConfig | None = OpenAIGPTConfig()
     segment_length: int = 1  # number of sentences per segment
     query: str = ""  # query for relevance extraction
-    system_message = """
+    system_message: str = """
     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 

langroid/agent/special/sql/sql_chat_agent.py

@@ -184,7 +184,7 @@ class SQLChatAgent(ChatAgent):
         if self.config.use_helper:
             # helper_config.system_message is now the fully-populated sys msg of
             # the main SQLAgent.
-            self.helper_config = self.config.copy()
+            self.helper_config = self.config.model_copy()
             self.helper_config.is_helper = True
             self.helper_config.use_helper = False
             self.helper_config.chat_mode = False
@@ -271,8 +271,16 @@ class SQLChatAgent(ChatAgent):
 
     def _init_tools(self) -> None:
         """Initialize sys msg and tools."""
-        RunQueryTool._max_retained_tokens = self.config.max_retained_tokens
-        self.enable_message([RunQueryTool, ForwardTool])
+        # Create a custom RunQueryTool class with the desired max_retained_tokens
+        if self.config.max_retained_tokens is not None:
+
+            class CustomRunQueryTool(RunQueryTool):
+                _max_retained_tokens = self.config.max_retained_tokens
+
+            self.enable_message([CustomRunQueryTool, ForwardTool])
+        else:
+            self.enable_message([RunQueryTool, ForwardTool])
+
         if self.config.use_schema_tools:
             self._enable_schema_tools()
         if not self.config.chat_mode: