langroid 0.20.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

langroid/agent/base.py

@@ -880,7 +880,13 @@ class Agent(ABC):
         return cdoc
 
     def has_tool_message_attempt(self, msg: str | ChatDocument | None) -> bool:
-        """Check whether msg contains a Tool/fn-call attempt (by the LLM)"""
+        """
+        Check whether msg contains a Tool/fn-call attempt (by the LLM).
+
+        CAUTION: This uses self.get_tool_messages(msg) which as a side-effect
+        may update msg.tool_messages when msg is a ChatDocument, if there are
+        any tools in msg.
+        """
         if msg is None:
             return False
         try:
@@ -921,6 +927,9 @@ class Agent(ABC):
     ) -> List[ToolMessage]:
         """
         Get ToolMessages recognized in msg, handle-able by this agent.
+        NOTE: as a side-effect, this will update msg.tool_messages
+        when msg is a ChatDocument and msg contains tool messages.
+
         If all_tools is True:
         - return all tools, i.e. any tool in self.llm_tools_known,
             whether it is handled by this agent or not;

langroid/agent/special/arangodb/arangodb_agent.py

@@ -27,6 +27,7 @@ from langroid.agent.special.arangodb.tools import (
     aql_retrieval_tool_name,
     arango_schema_tool_name,
 )
+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
@@ -88,11 +89,14 @@ class QueryResult(BaseModel):
 class ArangoChatAgentConfig(ChatAgentConfig):
     arango_settings: ArangoSettings = ArangoSettings()
     system_message: str = DEFAULT_ARANGO_CHAT_SYSTEM_MESSAGE
-    kg_schema: Optional[Dict[str, List[Dict[str, Any]]]] = None
+    kg_schema: str | Dict[str, List[Dict[str, Any]]] | None = None
     database_created: bool = False
-    use_schema_tools: bool = True
+    prepopulate_schema: bool = True
     use_functions_api: bool = True
+    max_num_results: int = 10  # how many results to return from AQL query
     max_result_tokens: int = 1000  # truncate long results to this many tokens
+    max_schema_fields: int = 500  # max fields to show in schema
+    max_tries: int = 10  # how many attempts to answer user question
     use_tools: bool = False
     schema_sample_pct: float = 0
     # whether the agent is used in a continuous chat with user,
@@ -103,16 +107,81 @@ class ArangoChatAgentConfig(ChatAgentConfig):
 
 class ArangoChatAgent(ChatAgent):
     def __init__(self, config: ArangoChatAgentConfig):
+        super().__init__(config)
         self.config: ArangoChatAgentConfig = config
+        self.init_state()
         self._validate_config()
         self._import_arango()
         self._initialize_db()
         self._init_tools_sys_message()
-        self.init_state()
 
     def init_state(self) -> None:
         super().init_state()
         self.current_retrieval_aql_query: str = ""
+        self.current_schema_params: ArangoSchemaTool = ArangoSchemaTool()
+        self.num_tries = 0  # how many attempts to answer user question
+
+    def user_response(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        response = super().user_response(msg)
+        if response is None:
+            return None
+        response_str = response.content if response is not None else ""
+        if response_str != "":
+            self.num_tries = 0  # reset number of tries if user responds
+        return response
+
+    def llm_response(
+        self, message: Optional[str | ChatDocument] = None
+    ) -> Optional[ChatDocument]:
+        if self.num_tries > self.config.max_tries:
+            if self.config.chat_mode:
+                return self.create_llm_response(
+                    content=f"""
+                    {self.config.addressing_prefix}User
+                    I give up, since I have exceeded the 
+                    maximum number of tries ({self.config.max_tries}).
+                    Feel free to give me some hints!
+                    """
+                )
+            else:
+                return self.create_llm_response(
+                    tool_messages=[
+                        DoneTool(
+                            content=f"""
+                            Exceeded maximum number of tries ({self.config.max_tries}).
+                            """
+                        )
+                    ]
+                )
+
+        if isinstance(message, ChatDocument) and message.metadata.sender == Entity.USER:
+            message.content = (
+                message.content
+                + "\n"
+                + """
+                (REMEMBER, Do NOT use more than ONE TOOL/FUNCTION at a time!
+                you must WAIT for a helper to send you the RESULT(S) before
+                making another TOOL/FUNCTION call)
+                """
+            )
+
+        response = super().llm_response(message)
+        if (
+            response is not None
+            and self.config.chat_mode
+            and self.config.addressing_prefix in response.content
+            and self.has_tool_message_attempt(response)
+        ):
+            # response contains both a user-addressing and a tool, which
+            # is not allowed, so remove the user-addressing prefix
+            response.content = response.content.replace(
+                self.config.addressing_prefix, ""
+            )
+
+        return response
 
     def _validate_config(self) -> None:
         assert isinstance(self.config, ArangoChatAgentConfig)
@@ -230,6 +299,7 @@ class ArangoChatAgent(ChatAgent):
             try:
                 cursor = self.db.aql.execute(query, bind_vars=bind_vars)
                 records = [doc for doc in cursor]  # type: ignore
+                records = records[: self.config.max_num_results]
                 logger.warning(f"Records retrieved: {records}")
                 return QueryResult(success=True, data=records if records else [])
             except Exception as e:
@@ -273,6 +343,28 @@ class ArangoChatAgent(ChatAgent):
                 success=False, data=f"Failed after max retries: {str(e)}"
             )
 
+    def _limit_tokens(self, text: str) -> str:
+        result = text
+        n_toks = self.num_tokens(result)
+        if n_toks > self.config.max_result_tokens:
+            logger.warning(
+                f"""
+                Your query resulted in a large result of  
+                {n_toks} tokens,
+                which will be truncated to {self.config.max_result_tokens} tokens.
+                If this does not give satisfactory results,
+                please retry with a more focused query.
+                """
+            )
+            if self.parser is not None:
+                result = self.parser.truncate_tokens(
+                    result,
+                    self.config.max_result_tokens,
+                )
+            else:
+                result = result[: self.config.max_result_tokens * 4]  # truncate roughly
+        return result
+
     def aql_retrieval_tool(self, msg: AQLRetrievalTool) -> str:
         """Handle AQL query for data retrieval"""
         if not self.tried_schema:
@@ -285,7 +377,13 @@ class ArangoChatAgent(ChatAgent):
             return """
             You need to create the database first using `{aql_creation_tool_name}`.
             """
+        self.num_tries += 1
         query = msg.aql_query
+        if query == self.current_retrieval_aql_query:
+            return """
+            You have already tried this query, so you will get the same results again!
+            If you need to retry, please MODIFY the query to get different results.
+            """
         self.current_retrieval_aql_query = query
         logger.info(f"Executing AQL query: {query}")
         response = self.read_query(query)
@@ -299,28 +397,11 @@ class ArangoChatAgent(ChatAgent):
             """
         # truncate long results
         result = str(response.data)
-        n_toks = self.num_tokens(result)
-        if n_toks > self.config.max_result_tokens:
-            logger.warning(
-                f"""
-                Your query resulted in a large result of  
-                {n_toks} tokens,
-                which will be truncated to {self.config.max_result_tokens} tokens.
-                If this does not give satisfactory results,
-                please retry with a more focused query.
-                """
-            )
-            if self.parser is not None:
-                result = self.parser.truncate_tokens(
-                    result,
-                    self.config.max_result_tokens,
-                )
-            else:
-                result = result[: self.config.max_result_tokens * 4]  # truncate roughly
-        return result
+        return self._limit_tokens(result)
 
     def aql_creation_tool(self, msg: AQLCreationTool) -> str:
         """Handle AQL query for creating data"""
+        self.num_tries += 1
         query = msg.aql_query
         logger.info(f"Executing AQL query: {query}")
         response = self.write_query(query)
@@ -334,12 +415,45 @@ class ArangoChatAgent(ChatAgent):
         self,
         msg: ArangoSchemaTool | None,
     ) -> Dict[str, List[Dict[str, Any]]] | str:
-        """Get database schema including collections, properties, and relationships"""
+        """Get database schema. If collections=None, include all collections.
+        If properties=False, show only connection info,
+        else show all properties and example-docs.
+        """
+
+        if (
+            msg is not None
+            and msg.collections == self.current_schema_params.collections
+            and msg.properties == self.current_schema_params.properties
+        ):
+            return """
+            You have already tried this schema TOOL, so you will get the same results 
+            again! Please MODIFY the tool params `collections` or `properties` to get
+            different results.
+            """
+
+        if msg is not None:
+            collections = msg.collections
+            properties = msg.properties
+        else:
+            collections = None
+            properties = True
         self.tried_schema = True
-        if self.config.kg_schema is not None and len(self.config.kg_schema) > 0:
+        if (
+            self.config.kg_schema is not None
+            and len(self.config.kg_schema) > 0
+            and msg is None
+        ):
+            # we are trying to pre-populate full schema before the agent runs,
+            # so get it if it's already available
+            # (Note of course that this "full schema" may actually be incomplete)
             return self.config.kg_schema
+
+        # increment tries only if the LLM is asking for the schema,
+        # in which case msg will not be None
+        self.num_tries += msg is not None
+
         try:
-            # Get graph schemas
+            # Get graph schemas (keeping full graph info)
             graph_schema = [
                 {"graph_name": g["name"], "edge_definitions": g["edge_definitions"]}
                 for g in self.db.graphs()  # type: ignore
@@ -348,57 +462,78 @@ class ArangoChatAgent(ChatAgent):
             # Get collection schemas
             collection_schema = []
             for collection in self.db.collections():  # type: ignore
-                if collection["name"].startswith("_"):  # Skip system collections
+                if collection["name"].startswith("_"):
                     continue
 
                 col_name = collection["name"]
+                if collections and col_name not in collections:
+                    continue
+
                 col_type = collection["type"]
                 col_size = self.db.collection(col_name).count()
 
-                if col_size == 0:  # Skip empty collections
+                if col_size == 0:
                     continue
 
-                # Calculate sample size
-                limit_amount = (
-                    ceil(
-                        self.config.schema_sample_pct * col_size / 100.0  # type: ignore
-                    )
-                    or 1
-                )
-
-                # Query to get sample documents and their properties
-                sample_query = f"""
-                    FOR doc in {col_name}
-                    LIMIT {limit_amount}
-                    RETURN doc
-                """
+                if properties:
+                    # Full property collection with sampling
+                    lim = self.config.schema_sample_pct * col_size  # type: ignore
+                    limit_amount = ceil(lim / 100.0) or 1
+                    sample_query = f"""
+                        FOR doc in {col_name}
+                        LIMIT {limit_amount}
+                        RETURN doc
+                    """
 
-                properties = []
-                example_doc = None
-
-                def simplify_doc(doc: Any) -> Any:
-                    if isinstance(doc, list) and len(doc) > 0:
-                        return [simplify_doc(doc[0])]
-                    if isinstance(doc, dict):
-                        return {k: simplify_doc(v) for k, v in doc.items()}
-                    return doc
-
-                for doc in self.db.aql.execute(sample_query):  # type: ignore
-                    if example_doc is None:
-                        example_doc = simplify_doc(doc)
-                    for key, value in doc.items():
-                        prop = {"name": key, "type": type(value).__name__}
-                        if prop not in properties:
-                            properties.append(prop)
-
-                collection_schema.append(
-                    {
+                    properties_list = []
+                    example_doc = None
+
+                    def simplify_doc(doc: Any) -> Any:
+                        if isinstance(doc, list) and len(doc) > 0:
+                            return [simplify_doc(doc[0])]
+                        if isinstance(doc, dict):
+                            return {k: simplify_doc(v) for k, v in doc.items()}
+                        return doc
+
+                    for doc in self.db.aql.execute(sample_query):  # type: ignore
+                        if example_doc is None:
+                            example_doc = simplify_doc(doc)
+                        for key, value in doc.items():
+                            prop = {"name": key, "type": type(value).__name__}
+                            if prop not in properties_list:
+                                properties_list.append(prop)
+
+                    collection_schema.append(
+                        {
+                            "collection_name": col_name,
+                            "collection_type": col_type,
+                            f"{col_type}_properties": properties_list,
+                            f"example_{col_type}": example_doc,
+                        }
+                    )
+                else:
+                    # Basic info + from/to for edges only
+                    collection_info = {
                         "collection_name": col_name,
                         "collection_type": col_type,
-                        f"{col_type}_properties": properties,
-                        f"example_{col_type}": example_doc,
                     }
-                )
+                    if col_type == "edge":
+                        # Get a sample edge to extract from/to fields
+                        sample_edge = next(
+                            self.db.aql.execute(  # type: ignore
+                                f"FOR e IN {col_name} LIMIT 1 RETURN e"
+                            ),
+                            None,
+                        )
+                        if sample_edge:
+                            collection_info["from_collection"] = sample_edge[
+                                "_from"
+                            ].split("/")[0]
+                            collection_info["to_collection"] = sample_edge["_to"].split(
+                                "/"
+                            )[0]
+
+                    collection_schema.append(collection_info)
 
             schema = {
                 "Graph Schema": graph_schema,
@@ -406,10 +541,41 @@ class ArangoChatAgent(ChatAgent):
             }
             schema_str = json.dumps(schema, indent=2)
             logger.warning(f"Schema retrieved:\n{schema_str}")
-            # save schema to file "logs/arangoo-schema.json"
             with open("logs/arango-schema.json", "w") as f:
                 f.write(schema_str)
-            self.config.kg_schema = schema  # type: ignore
+            if (n_fields := count_fields(schema)) > self.config.max_schema_fields:
+                logger.warning(
+                    f"""
+                    Schema has {n_fields} fields, which exceeds the maximum of
+                    {self.config.max_schema_fields}. Showing a trimmed version
+                    that only includes edge info and no other properties.
+                    """
+                )
+                schema = trim_schema(schema)
+                n_fields = count_fields(schema)
+                logger.warning(f"Schema trimmed down to {n_fields} fields.")
+                schema_str = (
+                    json.dumps(schema)
+                    + "\n"
+                    + f"""
+                    
+                    CAUTION: The requested schema was too large, so 
+                    the schema has been trimmed down to show only all collection names,
+                    their types, 
+                    and edge relationships (from/to collections) without any properties.
+                    To find out more about the schema, you can EITHER:
+                    - Use the `{arango_schema_tool_name}` tool again with the 
+                      `properties` arg set to True, and `collections` arg set to
+                        specific collections you want to know more about, OR
+                    - Use the `{aql_retrieval_tool_name}` tool to learn more about
+                      the schema by querying the database.
+                      
+                    """
+                )
+                if msg is None:
+                    self.config.kg_schema = schema_str
+                return schema_str
+            self.config.kg_schema = schema
             return schema
 
         except Exception as e:
@@ -432,9 +598,10 @@ class ArangoChatAgent(ChatAgent):
 
         super().__init__(self.config)
         # Note we are enabling GraphSchemaTool regardless of whether
-        # self.config.use_schema_tools is True or False, because
+        # self.config.prepopulate_schema is True or False, because
         # even when schema provided, the agent may later want to get the schema,
-        # e.g. if the db evolves, or if it needs to bring in the schema
+        # e.g. if the db evolves, or schema was trimmed due to size, or
+        # if it needs to bring in the schema into recent context.
 
         self.enable_message(
             [
@@ -454,7 +621,7 @@ class ArangoChatAgent(ChatAgent):
         assert isinstance(self.config, ArangoChatAgentConfig)
         return (
             SCHEMA_TOOLS_SYS_MSG
-            if self.config.use_schema_tools
+            if not self.config.prepopulate_schema
             else SCHEMA_PROVIDED_SYS_MSG.format(schema=self.arango_schema_tool(None))
         )
 

langroid/agent/special/arangodb/system_messages.py

@@ -9,7 +9,7 @@ done_tool_name = DoneTool.default_value("request")
 
 arango_schema_tool_description = f"""
 `{arango_schema_tool_name}` tool/function-call to find the schema
-of the graph database, i.e. get all the collections
+of the graph database, or for some SPECIFIC collections, i.e. get information on 
 (document and edge), their attributes, and graph definitions available in your
 ArangoDB database. You MUST use this tool BEFORE attempting to use the
 `{aql_retrieval_tool_name}` tool/function-call, to ensure that you are using the
@@ -18,7 +18,8 @@ correct collection names and attributes in your `{aql_retrieval_tool_name}` tool
 
 aql_retrieval_tool_description = f"""
 `{aql_retrieval_tool_name}` tool/function-call to retrieve information from 
-  the database using AQL (ArangoDB Query Language) queries.
+  the database using AQL (ArangoDB Query Language) queries, to answer
+  the user's questions, OR for you to learn more about the SCHEMA of the database.
 """
 
 aql_creation_tool_description = f"""
@@ -26,6 +27,29 @@ aql_creation_tool_description = f"""
 documents/edges in the database.
 """
 
+aql_retrieval_query_example = """
+EXAMPLE:
+Suppose you are asked this question "Does Bob have a father?".
+Then you will go through the following steps, where YOU indicates
+the message YOU will be sending, and RESULTS indicates the RESULTS
+you will receive from the helper executing the query:
+
+1. YOU:
+    {{ "request": "aql_retrieval_tool",
+      "aql_query": "FOR v, e, p in ... [query truncated for brevity]..."}}
+
+    2. RESULTS:
+    [.. results from the query...]
+    3. YOU: [ since results were not satisfactory, you try ANOTHER query]
+    {{ "request": "aql_retrieval_tool",
+    "aql_query": "blah blah ... [query truncated for brevity]..."}}
+    }}
+    4. RESULTS:
+    [.. results from the query...]
+    5. YOU: [ now you have the answer, you can generate your response ]
+    The answer is YES, Bob has a father, and his name is John.
+"""
+
 aql_query_instructions = """
 When writing AQL queries:
 1. Use the exact property names shown in the schema
@@ -63,6 +87,7 @@ REMEMBER:
     with your response. DO NOT MAKE UP RESULTS FROM A TOOL!
 [3] YOU MUST NOT ANSWER queries from your OWN KNOWLEDGE; ALWAYS RELY ON 
     the result of a TOOL/FUNCTION to compose your response.
+[4] Use ONLY ONE TOOL/FUNCTION at a TIME!
 """
 # sys msg to use when schema already provided initially,
 # so agent should not use schema tool
@@ -77,6 +102,7 @@ and their attribute keys available in your ArangoDB database.
 {{schema}}
 === END SCHEMA ===
 
+
 To help with the user's question or database update/creation request, 
 you have access to these tools:
 
@@ -84,10 +110,6 @@ you have access to these tools:
 
 - {aql_creation_tool_description}
 
-Since the schema has been provided, you may not need to use the tool below,
-but you may use it if you need to remind yourself about the schema:
-
-- {arango_schema_tool_description}
 
 {tool_result_instruction}
 """
@@ -113,27 +135,34 @@ DEFAULT_ARANGO_CHAT_SYSTEM_MESSAGE = f"""
 {{mode}}
 
 You do not need to be able to answer a question with just one query. 
-You could make a sequence of AQL queries to find the answer to the question.
+You can make a query, WAIT for the result, 
+THEN make ANOTHER query, WAIT for result,
+THEN make ANOTHER query, and so on, until you have the answer.
 
 {aql_query_instructions}
 
 RETRY-SUGGESTIONS:
 If you receive a null or other unexpected result,
 (a) make sure you use the available TOOLs correctly,
-(b) USE `{arango_schema_tool_name}` tool/function-call to get all collections, 
-    their attributes and graph definitions available in your ArangoDB database.
+(b) learn more about the schema using EITHER:
+ - `{arango_schema_tool_name}` tool/function-call to find properties of specific
+    collections or other parts of the schema, OR
+ - `{aql_retrieval_tool_name}` tool/function-call to use AQL queries to 
+    find specific parts of the schema.
 (c) Collection names are CASE-SENSITIVE -- make sure you adhere to the exact 
     collection name you found in the schema.
 (d) see if you have made an assumption in your AQL query, and try another way, 
     or use `{aql_retrieval_tool_name}` to explore the database contents before 
     submitting your final query. 
-(f) Try APPROXIMATE or PARTIAL MATCHES to strings in the user's query, 
+(e) Try APPROXIMATE or PARTIAL MATCHES to strings in the user's query, 
     e.g. user may ask about "Godfather" instead of "The Godfather",
     or try using CASE-INSENSITIVE MATCHES.
     
 Start by asking what the user needs help with.
 
 {tool_result_instruction}
+
+{aql_retrieval_query_example}
 """
 
 ADDRESSING_INSTRUCTION = """

langroid/agent/special/arangodb/tools.py

@@ -1,3 +1,5 @@
+from typing import List, Tuple
+
 from langroid.agent.tool_message import ToolMessage
 
 
@@ -5,11 +7,46 @@ class AQLRetrievalTool(ToolMessage):
     request: str = "aql_retrieval_tool"
     purpose: str = """
         To send an <aql_query> in response to a user's request/question, 
+        OR to find SCHEMA information,
         and WAIT for results of the <aql_query> BEFORE continuing with response.
         You will receive RESULTS from this tool, and ONLY THEN you can continue.
     """
     aql_query: str
 
+    @classmethod
+    def examples(cls) -> List[ToolMessage | Tuple[str, ToolMessage]]:
+        """Few-shot examples to include in tool instructions."""
+        return [
+            (
+                "I want to see who Bob's Father is",
+                cls(
+                    aql_query="""
+                    FOR v, e, p IN 1..1 OUTBOUND 'users/Bob' GRAPH 'family_tree'
+                    FILTER p.edges[0].type == 'father'
+                    RETURN v
+                    """
+                ),
+            ),
+            (
+                "I want to know the properties of the Actor node",
+                cls(
+                    aql_query="""
+                    FOR doc IN Actor
+                    LIMIT 1
+                    RETURN ATTRIBUTES(doc)                    
+                    """
+                ),
+            ),
+        ]
+
+    @classmethod
+    def instructions(cls) -> str:
+        return """
+        When using this TOOL/Function-call, you must WAIT to receive the RESULTS 
+        of the AQL query, before continuing your response!
+        DO NOT ASSUME YOU KNOW THE RESULTs BEFORE RECEIVING THEM.        
+        """
+
 
 aql_retrieval_tool_name = AQLRetrievalTool.default_value("request")
 
@@ -23,6 +60,23 @@ class AQLCreationTool(ToolMessage):
     """
     aql_query: str
 
+    @classmethod
+    def examples(cls) -> List[ToolMessage | Tuple[str, ToolMessage]]:
+        """Few-shot examples to include in tool instructions."""
+        return [
+            (
+                "Create a new document in the collection 'users'",
+                cls(
+                    aql_query="""
+                    INSERT {
+                      "name": "Alice",
+                      "age": 30
+                    } INTO users
+                    """
+                ),
+            ),
+        ]
+
 
 aql_creation_tool_name = AQLCreationTool.default_value("request")
 
@@ -30,10 +84,19 @@ aql_creation_tool_name = AQLCreationTool.default_value("request")
 class ArangoSchemaTool(ToolMessage):
     request: str = "arango_schema_tool"
     purpose: str = """
-        To get the schema of the Arango graph database.
+        To get the schema of the Arango graph database,
+        or some part of it. Follow these instructions:
+        1. Set <properties> to True to get the properties of the collections,
+        and False if you only want to see the graph structure and get only the
+        from/to relations of the edges.
+        2. Set <collections> to a list of collection names if you want to see,
+        or leave it as None to see all ALL collections.
         IMPORTANT: YOU MUST WAIT FOR THE RESULT OF THE TOOL BEFORE CONTINUING.
         You will receive RESULTS from this tool, and ONLY THEN you can continue.
     """
 
+    properties: bool = True
+    collections: List[str] | None = None
+
 
 arango_schema_tool_name = ArangoSchemaTool.default_value("request")

langroid/agent/special/arangodb/utils.py

@@ -0,0 +1,36 @@
+from typing import Any, Dict, List
+
+
+def count_fields(schema: Dict[str, List[Dict[str, Any]]]) -> int:
+    total = 0
+    for coll in schema["Collection Schema"]:
+        # Count all keys in each collection's dict
+        total += len(coll)
+        # Also count properties if they exist
+        props = coll.get(f"{coll['collection_type']}_properties", [])
+        total += len(props)
+    return total
+
+
+def trim_schema(
+    schema: Dict[str, List[Dict[str, Any]]]
+) -> Dict[str, List[Dict[str, Any]]]:
+    """Keep only edge connection info, remove properties and examples"""
+    trimmed: Dict[str, List[Dict[str, Any]]] = {
+        "Graph Schema": schema["Graph Schema"],
+        "Collection Schema": [],
+    }
+    for coll in schema["Collection Schema"]:
+        col_info: Dict[str, Any] = {
+            "collection_name": coll["collection_name"],
+            "collection_type": coll["collection_type"],
+        }
+        if coll["collection_type"] == "edge":
+            # preserve from/to info if present
+            if f"example_{coll['collection_type']}" in coll:
+                example = coll[f"example_{coll['collection_type']}"]
+                if example and "_from" in example:
+                    col_info["from_collection"] = example["_from"].split("/")[0]
+                    col_info["to_collection"] = example["_to"].split("/")[0]
+        trimmed["Collection Schema"].append(col_info)
+    return trimmed

langroid/agent/task.py

@@ -1330,6 +1330,9 @@ class Task:
                     max_cost=self.max_cost,
                     max_tokens=self.max_tokens,
                 )
+                # update result.tool_messages if any
+                if isinstance(result, ChatDocument):
+                    self.agent.get_tool_messages(result)
                 if result is not None:
                     content, id2result, oai_tool_id = self.agent.process_tool_results(
                         result.content,
@@ -1358,6 +1361,9 @@ class Task:
         else:
             response_fn = self._entity_responder_map[cast(Entity, e)]
             result = response_fn(self.pending_message)
+            # update result.tool_messages if any
+            if isinstance(result, ChatDocument):
+                self.agent.get_tool_messages(result)
 
         result_chat_doc = self.agent.to_ChatDocument(
             result,
@@ -1388,7 +1394,7 @@ class Task:
             # ignore all string-based signaling/routing
             return result
         # parse various routing/addressing strings in result
-        is_pass, recipient, content = parse_routing(
+        is_pass, recipient, content = self._parse_routing(
             result,
             addressing_prefix=self.config.addressing_prefix,
         )
@@ -1521,7 +1527,7 @@ class Task:
         oai_tool_id2result = result_msg.oai_tool_id2result if result_msg else None
         fun_call = result_msg.function_call if result_msg else None
         tool_messages = result_msg.tool_messages if result_msg else []
-        # if there is an LLMDoneTool or AgentDoneTool among these,
+        # if there is an DoneTool or AgentDoneTool among these,
         # we extract content and tools from here, and ignore all others
         for t in tool_messages:
             if isinstance(t, FinalResultTool):
@@ -1533,6 +1539,8 @@ class Task:
                 # there shouldn't be multiple tools like this; just take the first
                 content = to_string(t.content)
                 content_any = t.content
+                fun_call = None
+                oai_tool_calls = None
                 if isinstance(t, AgentDoneTool):
                     # AgentDoneTool may have tools, unlike DoneTool
                     tool_messages = t.tools
@@ -1940,58 +1948,72 @@ class Task:
         """
         self.color_log = enable
 
+    def _parse_routing(
+        self,
+        msg: ChatDocument | str,
+        addressing_prefix: str = "",
+    ) -> Tuple[bool | None, str | None, str | None]:
+        """
+        Parse routing instruction if any, of the form:
+        PASS:<recipient>  (pass current pending msg to recipient)
+        SEND:<recipient> <content> (send content to recipient)
+        @<recipient> <content> (send content to recipient)
+        Args:
+            msg (ChatDocument|str|None): message to parse
+            addressing_prefix (str): prefix to address other agents or entities,
+                 (e.g. "@". See documentation of `TaskConfig` for details).
+        Returns:
+            Tuple[bool|None, str|None, str|None]:
+                bool: true=PASS, false=SEND, or None if neither
+                str: recipient, or None
+                str: content to send, or None
+        """
+        # handle routing instruction-strings in result if any,
+        # such as PASS, PASS_TO, or SEND
 
-def parse_routing(
-    msg: ChatDocument | str,
-    addressing_prefix: str = "",
-) -> Tuple[bool | None, str | None, str | None]:
-    """
-    Parse routing instruction if any, of the form:
-    PASS:<recipient>  (pass current pending msg to recipient)
-    SEND:<recipient> <content> (send content to recipient)
-    @<recipient> <content> (send content to recipient)
-    Args:
-        msg (ChatDocument|str|None): message to parse
-        addressing_prefix (str): prefix to address other agents or entities,
-             (e.g. "@". See documentation of `TaskConfig` for details).
-    Returns:
-        Tuple[bool|None, str|None, str|None]:
-            bool: true=PASS, false=SEND, or None if neither
-            str: recipient, or None
-            str: content to send, or None
-    """
-    # handle routing instruction in result if any,
-    # of the form PASS=<recipient>
-    content = msg.content if isinstance(msg, ChatDocument) else msg
-    content = content.strip()
-    if PASS in content and PASS_TO not in content:
-        return True, None, None
-    if PASS_TO in content and content.split(":")[1] != "":
-        return True, content.split(":")[1], None
-    if (
-        SEND_TO in content
-        and (addressee_content := parse_addressed_message(content, SEND_TO))[0]
-        is not None
-    ):
-        (addressee, content_to_send) = addressee_content
-        # if no content then treat same as PASS_TO
-        if content_to_send == "":
-            return True, addressee, None
-        else:
-            return False, addressee, content_to_send
-    if (
-        addressing_prefix != ""
-        and addressing_prefix in content
-        and (addressee_content := parse_addressed_message(content, addressing_prefix))[
-            0
-        ]
-        is not None
-    ):
-        (addressee, content_to_send) = addressee_content
-        # if no content then treat same as PASS_TO
-        if content_to_send == "":
-            return True, addressee, None
-        else:
-            return False, addressee, content_to_send
+        msg_str = msg.content if isinstance(msg, ChatDocument) else msg
+        if (
+            self.agent.has_tool_message_attempt(msg)
+            and not msg_str.startswith(PASS)
+            and not msg_str.startswith(PASS_TO)
+            and not msg_str.startswith(SEND_TO)
+        ):
+            # if there's an attempted tool-call, we ignore any routing strings,
+            # unless they are at the start of the msg
+            return None, None, None
+
+        content = msg.content if isinstance(msg, ChatDocument) else msg
+        content = content.strip()
+        if PASS in content and PASS_TO not in content:
+            return True, None, None
+        if PASS_TO in content and content.split(":")[1] != "":
+            return True, content.split(":")[1], None
+        if (
+            SEND_TO in content
+            and (addressee_content := parse_addressed_message(content, SEND_TO))[0]
+            is not None
+        ):
+            # Note this will discard any portion of content BEFORE SEND_TO.
+            # TODO maybe make this configurable.
+            (addressee, content_to_send) = addressee_content
+            # if no content then treat same as PASS_TO
+            if content_to_send == "":
+                return True, addressee, None
+            else:
+                return False, addressee, content_to_send
+        if (
+            addressing_prefix != ""
+            and addressing_prefix in content
+            and (
+                addressee_content := parse_addressed_message(content, addressing_prefix)
+            )[0]
+            is not None
+        ):
+            (addressee, content_to_send) = addressee_content
+            # if no content then treat same as PASS_TO
+            if content_to_send == "":
+                return True, addressee, None
+            else:
+                return False, addressee, content_to_send
 
-    return None, None, None
+        return None, None, None

langroid/agent/tools/orchestration.py

@@ -41,11 +41,11 @@ class DoneTool(ToolMessage):
     """Tool for Agent Entity (i.e. agent_response) or LLM entity (i.e. llm_response) to
     signal the current task is done, with some content as the result."""
 
-    purpose = """
+    purpose: str = """
     To signal the current task is done, along with an optional message <content>
     of arbitrary type (default None).
     """
-    request = "done_tool"
+    request: str = "done_tool"
     content: str = ""
 
     def response(self, agent: ChatAgent) -> ChatDocument:
@@ -77,7 +77,7 @@ class ResultTool(ToolMessage):
     Note:
         - when defining a tool handler or agent_response, you can directly return
             ResultTool(field1 = val1, ...),
-            where the values can be aribitrary data structures, including nested
+            where the values can be arbitrary data structures, including nested
             Pydantic objs, or you can define a subclass of ResultTool with the
             fields you want to return.
         - This is a special ToolMessage that is NOT meant to be used or handled
@@ -143,10 +143,10 @@ class PassTool(ToolMessage):
     Similar to ForwardTool, but without specifying the recipient agent.
     """
 
-    purpose = """
+    purpose: str = """
     To pass the current message so that other agents can handle it.
     """
-    request = "pass_tool"
+    request: str = "pass_tool"
 
     def response(self, agent: ChatAgent, chat_doc: ChatDocument) -> ChatDocument:
         """When this tool is enabled for an Agent, this will result in a method
@@ -178,10 +178,10 @@ class DonePassTool(PassTool):
     Similar to PassTool, except we append a DoneTool to the result tool_messages.
     """
 
-    purpose = """
+    purpose: str = """
     To signal the current task is done, with results set to the current/incoming msg.
     """
-    request = "done_pass_tool"
+    request: str = "done_pass_tool"
 
     def response(self, agent: ChatAgent, chat_doc: ChatDocument) -> ChatDocument:
         # use PassTool to get the right ChatDocument to pass...

langroid/language_models/__init__.py

@@ -17,6 +17,7 @@ from .base import (
 from .openai_gpt import (
     OpenAIChatModel,
     AnthropicModel,
+    GeminiModel,
     OpenAICompletionModel,
     OpenAIGPTConfig,
     OpenAIGPT,
@@ -41,6 +42,7 @@ __all__ = [
     "LLMResponse",
     "OpenAIChatModel",
     "AnthropicModel",
+    "GeminiModel",
     "OpenAICompletionModel",
     "OpenAIGPTConfig",
     "OpenAIGPT",

langroid/language_models/openai_gpt.py

@@ -66,6 +66,7 @@ if "OLLAMA_HOST" in os.environ:
 else:
     OLLAMA_BASE_URL = "http://localhost:11434/v1"
 
+GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/"
 OLLAMA_API_KEY = "ollama"
 DUMMY_API_KEY = "xxx"
 
@@ -92,6 +93,14 @@ class OpenAIChatModel(str, Enum):
     O1_MINI = "o1-mini"
 
 
+class GeminiModel(str, Enum):
+    """Enum for Gemini models"""
+
+    GEMINI_1_5_FLASH = "gemini-1.5-flash"
+    GEMINI_1_5_FLASH_8B = "gemini-1.5-flash-8b"
+    GEMINI_1_5_PRO = "gemini-1.5-pro"
+
+
 class OpenAICompletionModel(str, Enum):
     """Enum for OpenAI Completion models"""
 
@@ -503,6 +512,7 @@ class OpenAIGPT(LanguageModel):
 
         self.is_groq = self.config.chat_model.startswith("groq/")
         self.is_cerebras = self.config.chat_model.startswith("cerebras/")
+        self.is_gemini = self.config.chat_model.startswith("gemini/")
 
         if self.is_groq:
             self.config.chat_model = self.config.chat_model.replace("groq/", "")
@@ -524,6 +534,11 @@ class OpenAIGPT(LanguageModel):
                 api_key=self.api_key,
             )
         else:
+            if self.is_gemini:
+                self.config.chat_model = self.config.chat_model.replace("gemini/", "")
+                self.api_key = os.getenv("GEMINI_API_KEY", DUMMY_API_KEY)
+                self.api_base = GEMINI_BASE_URL
+
             self.client = OpenAI(
                 api_key=self.api_key,
                 base_url=self.api_base,
@@ -623,7 +638,13 @@ class OpenAIGPT(LanguageModel):
         Currently main troublemaker is o1* series.
         """
         match self.config.chat_model:
-            case OpenAIChatModel.O1_MINI | OpenAIChatModel.O1_PREVIEW:
+            case (
+                OpenAIChatModel.O1_MINI
+                | OpenAIChatModel.O1_PREVIEW
+                | GeminiModel.GEMINI_1_5_FLASH
+                | GeminiModel.GEMINI_1_5_FLASH_8B
+                | GeminiModel.GEMINI_1_5_PRO
+            ):
                 return {"max_tokens": "max_completion_tokens"}
             case _:
                 return {}
@@ -1229,13 +1250,10 @@ class OpenAIGPT(LanguageModel):
         cost = 0.0
         prompt_tokens = 0
         completion_tokens = 0
-        if not cached and not self.get_stream():
+        if not cached and not self.get_stream() and response["usage"] is not None:
             prompt_tokens = response["usage"]["prompt_tokens"]
             completion_tokens = response["usage"]["completion_tokens"]
-            cost = self._cost_chat_model(
-                response["usage"]["prompt_tokens"],
-                response["usage"]["completion_tokens"],
-            )
+            cost = self._cost_chat_model(prompt_tokens, completion_tokens)
 
         return LLMTokenUsage(
             prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, cost=cost

{langroid-0.20.0.dist-info → langroid-0.21.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.20.0
+Version: 0.21.0
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani
@@ -248,6 +248,9 @@ teacher_task.run()
 <details>
 <summary> <b>Click to expand</b></summary>
 
+- **Nov 2024:**
+  - **[0.20.0](https://github.com/langroid/langroid/releases/tag/0.20.0)** Support for 
+    ArangoDB Knowledge Graphs.
 - **Oct 2024:**
   - **[0.18.0]** [LLMConfig.async_stream_quiet](https://langroid.github.io/langroid/notes/async-streaming/) flag to 
     turn off LLM output in async + stream mode.

{langroid-0.20.0.dist-info → langroid-0.21.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 langroid/__init__.py,sha256=z_fCOLQJPOw3LLRPBlFB5-2HyCjpPgQa4m4iY5Fvb8Y,1800
 langroid/agent/__init__.py,sha256=ll0Cubd2DZ-fsCMl7e10hf9ZjFGKzphfBco396IKITY,786
-langroid/agent/base.py,sha256=uCjFJ0Mjm6hu4MSCsz7vzGwzpbkrxddDIduhxS24jps,65034
+langroid/agent/base.py,sha256=sOZapdzHaB4kbCLu8vI_zZx78jIhv9fmWn0EWV4yTAE,65371
 langroid/agent/batch.py,sha256=QZdlt1563hx4l3AXrCaGovE-PNG93M3DsvQAbDzdiS8,13705
 langroid/agent/callbacks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/callbacks/chainlit.py,sha256=JJXI3UGTyTDg2FFath4rqY1GyUo_0pbVBt8CZpvdtn4,23289
@@ -11,9 +11,10 @@ langroid/agent/junk,sha256=LxfuuW7Cijsg0szAzT81OjWWv1PMNI-6w_-DspVIO2s,339
 langroid/agent/openai_assistant.py,sha256=2rjCZw45ysNBEGNzQM4uf0bTC4KkatGYAWcVcW4xcek,34337
 langroid/agent/special/__init__.py,sha256=gik_Xtm_zV7U9s30Mn8UX3Gyuy4jTjQe9zjiE3HWmEo,1273
 langroid/agent/special/arangodb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langroid/agent/special/arangodb/arangodb_agent.py,sha256=5X8yQr1MdTLchNK6L7g3gArFYVEiql4g35YPMw_4TIU,19579
-langroid/agent/special/arangodb/system_messages.py,sha256=j-vNz4SLKzE6-1a31ZFw7NHMD_uBepS8vxJ_bsjgMC4,5801
-langroid/agent/special/arangodb/tools.py,sha256=3CW2HsLqcM0JGemHsmihm8F5zZwwbFvAwyHzzaJ2S90,1319
+langroid/agent/special/arangodb/arangodb_agent.py,sha256=CXD1Z4yXnWwre4ybg7BMhVFTgQq0zVXzobutyxIYCP4,26862
+langroid/agent/special/arangodb/system_messages.py,sha256=udwfLleTdyz_DuxHuoiv2wHEZoAPBPbwdF_ivjIfP5c,6867
+langroid/agent/special/arangodb/tools.py,sha256=WasFERC1cToLOWi1cWqUs-TujU0A68gZWbhbP128obo,3499
+langroid/agent/special/arangodb/utils.py,sha256=LIevtkayIdVVXyj3jlbKH2WgdZTtH5-JLgbXOHC7uxs,1420
 langroid/agent/special/doc_chat_agent.py,sha256=xIqBOyLax_jMU0UevxqXf_aQUrRkW6MQUKpKnKvaqkQ,59281
 langroid/agent/special/lance_doc_chat_agent.py,sha256=s8xoRs0gGaFtDYFUSIRchsgDVbS5Q3C2b2mr3V1Fd-Q,10419
 langroid/agent/special/lance_rag/__init__.py,sha256=QTbs0IVE2ZgDg8JJy1zN97rUUg4uEPH7SLGctFNumk4,174
@@ -37,14 +38,14 @@ langroid/agent/special/sql/utils/system_message.py,sha256=qKLHkvQWRQodTtPLPxr1GS
 langroid/agent/special/sql/utils/tools.py,sha256=vFYysk6Vi7HJjII8B4RitA3pt_z3gkSglDNdhNVMiFc,1332
 langroid/agent/special/table_chat_agent.py,sha256=d9v2wsblaRx7oMnKhLV7uO_ujvk9gh59pSGvBXyeyNc,9659
 langroid/agent/structured_message.py,sha256=y7pud1EgRNeTFZlJmBkLmwME3yQJ_IYik-Xds9kdZbY,282
-langroid/agent/task.py,sha256=p7YsccpRIZcX0dFo5ll9EAaetWUcxZnCZal0i7Xsl1Y,85433
+langroid/agent/task.py,sha256=f7clh6p6Md0G4YGHqbFeeT88U4XoP0i3eatekV21hHE,86643
 langroid/agent/tool_message.py,sha256=jkN7uq7YwUC_wBcSCNUYjrB_His2YCfQay_lqIa4Tww,10498
 langroid/agent/tools/__init__.py,sha256=IMgCte-_ZIvCkozGQmvMqxIw7_nKLKzD78ccJL1bnQU,804
 langroid/agent/tools/duckduckgo_search_tool.py,sha256=NhsCaGZkdv28nja7yveAhSK_w6l_Ftym8agbrdzqgfo,1935
 langroid/agent/tools/file_tools.py,sha256=GjPB5YDILucYapElnvvoYpGJuZQ25ecLs2REv7edPEo,7292
 langroid/agent/tools/google_search_tool.py,sha256=y7b-3FtgXf0lfF4AYxrZ3K5pH2dhidvibUOAGBE--WI,1456
 langroid/agent/tools/metaphor_search_tool.py,sha256=qj4gt453cLEX3EGW7nVzVu6X7LCdrwjSlcNY0qJW104,2489
-langroid/agent/tools/orchestration.py,sha256=vp2Qx-DYPtDnACosxKqwHGy6DeD1QnEllWz0Ht81Cyc,10880
+langroid/agent/tools/orchestration.py,sha256=EDv1EMVGYqX82x3bCRbTn9gFNs66oosiUM8WTSZkUJg,10909
 langroid/agent/tools/recipient_tool.py,sha256=dr0yTxgNEIoxUYxH6TtaExC4G_8WdJ0xGohIa4dFLhY,9808
 langroid/agent/tools/retrieval_tool.py,sha256=2q2pfoYbZNfbWQ0McxrtmfF0ekGglIgRl-6uF26pa-E,871
 langroid/agent/tools/rewind_tool.py,sha256=XAXL3BpNhCmBGYq_qi_sZfHJuIw7NY2jp4wnojJ7WRs,5606
@@ -69,12 +70,12 @@ langroid/embedding_models/remote_embeds.py,sha256=6_kjXByVbqhY9cGwl9R83ZcYC2km-n
 langroid/exceptions.py,sha256=G60UVDChkUlBDVWHFr_43zUUszZHSejoU00tX_dfD68,2322
 langroid/language_models/.chainlit/config.toml,sha256=1t5lHORGzc2E6dkaO9P15jYHu2w-4Kl9pYjpDPc84vs,3716
 langroid/language_models/.chainlit/translations/en-US.json,sha256=DAFz2HjOFFfboCStrUfKFg2BpplJPK_OOtixwF_GivY,9931
-langroid/language_models/__init__.py,sha256=1sUGobooTqq77XC7LxKsvME0RgSd5GGmeyrPo9SMh4U,940
+langroid/language_models/__init__.py,sha256=8o8D8Lxaq961_oxVpB_bC2iEJ1GXJqYXMlwUcn6OJb8,976
 langroid/language_models/azure_openai.py,sha256=G4le3j4YLHV7IwgB2C37hO3MKijZ1KjynbYlEvpIF7Y,6214
 langroid/language_models/base.py,sha256=xMFg8syIHiA7ABRNWPXeFM1vGeY_1EN84ki8C3dycfw,22722
 langroid/language_models/config.py,sha256=9Q8wk5a7RQr8LGMT_0WkpjY8S4ywK06SalVRjXlfCiI,378
 langroid/language_models/mock_lm.py,sha256=HuiAvjHiCfffYF5xjFJUq945HVTW0QPbeUUctOnNCzQ,3868
-langroid/language_models/openai_gpt.py,sha256=-A2RB7bhzGUDUVM-JcBYO84WYsS2YHtb5xXqqM_dXG4,68874
+langroid/language_models/openai_gpt.py,sha256=CVRwYvIT8lt02GqI0G5_sTQXLcQbiwwNXpU8_qsw6q8,69584
 langroid/language_models/prompt_formatter/__init__.py,sha256=2-5cdE24XoFDhifOLl8yiscohil1ogbP1ECkYdBlBsk,372
 langroid/language_models/prompt_formatter/base.py,sha256=eDS1sgRNZVnoajwV_ZIha6cba5Dt8xjgzdRbPITwx3Q,1221
 langroid/language_models/prompt_formatter/hf_formatter.py,sha256=PVJppmjRvD-2DF-XNC6mE05vTZ9wbu37SmXwZBQhad0,5055
@@ -141,8 +142,8 @@ langroid/vector_store/meilisearch.py,sha256=6frB7GFWeWmeKzRfLZIvzRjllniZ1cYj3Hmh
 langroid/vector_store/momento.py,sha256=qR-zBF1RKVHQZPZQYW_7g-XpTwr46p8HJuYPCkfJbM4,10534
 langroid/vector_store/qdrant_cloud.py,sha256=3im4Mip0QXLkR6wiqVsjV1QvhSElfxdFSuDKddBDQ-4,188
 langroid/vector_store/qdrantdb.py,sha256=v88lqFkepADvlN6lByUj9I4NEKa9X9lWH16uTPPbYrE,17457
-pyproject.toml,sha256=DQbq0Xo4Bc3sJYF2-1qA9qnnKT0X-x5BzC99OrxcTNg,7488
-langroid-0.20.0.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
-langroid-0.20.0.dist-info/METADATA,sha256=5uni8xDGH809QGehHz0YfMguGFtvKDSJ9A7CfbduDnE,56758
-langroid-0.20.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-langroid-0.20.0.dist-info/RECORD,,
+pyproject.toml,sha256=yiLcLTge7QxJ1EWYE05yUHpBADouFh36WO3SlRzJKL0,7488
+langroid-0.21.0.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
+langroid-0.21.0.dist-info/METADATA,sha256=TSl0anzCbcnchyI2yTczV4fK2KBkfJnUu2Jl8oqGGLA,56893
+langroid-0.21.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+langroid-0.21.0.dist-info/RECORD,,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langroid"
-version = "0.20.0"
+version = "0.21.0"
 description = "Harness LLMs with Multi-Agent Programming"
 authors = ["Prasad Chalasani <pchalasani@gmail.com>"]
 readme = "README.md"