langroid 0.19.5__py3-none-any.whl → 0.20.1__py3-none-any.whl

langroid/agent/special/arangodb/system_messages.py

@@ -0,0 +1,183 @@
+from langroid.agent.special.arangodb.tools import (
+    aql_creation_tool_name,
+    aql_retrieval_tool_name,
+    arango_schema_tool_name,
+)
+from langroid.agent.tools.orchestration import DoneTool
+
+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, 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
+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, to answer
+  the user's questions, OR for you to learn more about the SCHEMA of the database.
+"""
+
+aql_creation_tool_description = f"""
+`{aql_creation_tool_name}` tool/function-call to execute AQL query that creates
+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
+2. Pay attention to the 'type' field of each node
+3. Note that all names are case-sensitive:
+   - collection names
+   - property names
+   - node type values
+   - relationship type values
+4. Always include type filters in your queries, e.g.:
+   FILTER doc.type == '<type-from-schema>'
+
+The schema shows:
+- Collections (usually 'nodes' and 'edges')
+- Node types in each collection
+- Available properties for each node type
+- Relationship types and their properties
+
+Examine the schema carefully before writing queries to ensure:
+- Correct property names
+- Correct node types
+- Correct relationship types
+
+You must be smart about using the right collection names and attributes
+based on the English description. If you are thinking of using a collection
+or attribute that does not exist, you are probably on the wrong track,
+so you should try your best to answer based on existing collections and attributes.
+DO NOT assume any collections or graphs other than those above.
+"""
+
+tool_result_instruction = """
+REMEMBER:
+[1]  DO NOT FORGET TO USE ONE OF THE AVAILABLE TOOLS TO ANSWER THE USER'S QUERY!!
+[2] When using a TOOL/FUNCTION, you MUST WAIT for the tool result before continuing
+    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
+SCHEMA_PROVIDED_SYS_MSG = f"""You are a data scientist and expert in Graph Databases, 
+with expertise in answering questions by interacting with an ArangoDB database.
+
+The schema below describes the ArangoDB database structure, 
+collections (document and edge),
+and their attribute keys available in your ArangoDB database.
+
+=== SCHEMA ===
+{{schema}}
+=== END SCHEMA ===
+
+
+To help with the user's question or database update/creation request, 
+you have access to these tools:
+
+- {aql_retrieval_tool_description}
+
+- {aql_creation_tool_description}
+
+
+{tool_result_instruction}
+"""
+
+# sys msg to use when schema is not initially provided,
+# and we want agent to use schema tool to get schema
+SCHEMA_TOOLS_SYS_MSG = f"""You are a data scientist and expert in 
+Arango Graph Databases, 
+with expertise in answering questions by querying ArangoDB database
+using the Arango Query Language (AQL).
+You have access to the following tools:
+
+- {arango_schema_tool_description}
+
+- {aql_retrieval_tool_description}
+
+- {aql_creation_tool_description}
+
+{tool_result_instruction}
+"""
+
+DEFAULT_ARANGO_CHAT_SYSTEM_MESSAGE = f"""
+{{mode}}
+
+You do not need to be able to answer a question with just one query. 
+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.
+(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.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 = """
+IMPORTANT - Whenever you are NOT writing an AQL query, make sure you address the 
+user using {prefix}User. You MUST use the EXACT syntax {prefix} !!!
+
+In other words, you ALWAYS EITHER:
+ - write an AQL query using one of the tools, 
+ - OR address the user using {prefix}User.
+ 
+YOU CANNOT ADDRESS THE USER WHEN USING A TOOL!!
+"""
+
+DONE_INSTRUCTION = f"""
+When you are SURE you have the CORRECT answer to a user's query or request, 
+use the `{done_tool_name}` with `content` set to the answer or result.
+If you DO NOT think you have the answer to the user's query or request,
+you SHOULD NOT use the `{done_tool_name}` tool.
+Instead, you must CONTINUE to improve your queries (tools) to get the correct answer,
+and finally use the `{done_tool_name}` tool to send the correct answer to the user.
+"""

langroid/agent/special/arangodb/tools.py

@@ -0,0 +1,102 @@
+from typing import List, Tuple
+
+from langroid.agent.tool_message import ToolMessage
+
+
+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")
+
+
+class AQLCreationTool(ToolMessage):
+    request: str = "aql_creation_tool"
+    purpose: str = """
+        To send the <aql_query> to create documents/edges in the graph database.
+        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.
+    """
+    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")
+
+
+class ArangoSchemaTool(ToolMessage):
+    request: str = "arango_schema_tool"
+    purpose: str = """
+        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/special/neo4j/neo4j_chat_agent.py

@@ -1,11 +1,9 @@
-import json
 import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from rich import print
 from rich.console import Console
 
-from langroid.agent import ToolMessage
 from langroid.pydantic_v1 import BaseModel, BaseSettings
 
 if TYPE_CHECKING:
@@ -13,13 +11,25 @@ if TYPE_CHECKING:
 
 from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
 from langroid.agent.chat_document import ChatDocument
-from langroid.agent.special.neo4j.utils.system_message import (
+from langroid.agent.special.neo4j.system_messages import (
     ADDRESSING_INSTRUCTION,
     DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE,
-    DEFAULT_SYS_MSG,
+    DONE_INSTRUCTION,
+    SCHEMA_PROVIDED_SYS_MSG,
     SCHEMA_TOOLS_SYS_MSG,
 )
+from langroid.agent.special.neo4j.tools import (
+    CypherCreationTool,
+    CypherRetrievalTool,
+    GraphSchemaTool,
+    cypher_creation_tool_name,
+    cypher_retrieval_tool_name,
+    graph_schema_tool_name,
+)
+from langroid.agent.tools.orchestration import DoneTool, ForwardTool
+from langroid.exceptions import LangroidImportError
 from langroid.mytypes import Entity
+from langroid.utils.constants import SEND_TO
 
 logger = logging.getLogger(__name__)
 
@@ -31,25 +41,6 @@ NEO4J_ERROR_MSG = "There was an error in your Cypher Query"
 # TOOLS to be used by the agent
 
 
-class CypherRetrievalTool(ToolMessage):
-    request: str = "retrieval_query"
-    purpose: str = """Use this tool to send the Cypher query to retreive data from the 
-    graph database based provided text description and schema."""
-    cypher_query: str
-
-
-class CypherCreationTool(ToolMessage):
-    request: str = "create_query"
-    purpose: str = """Use this tool to send the Cypher query to create 
-    entities/relationships in the graph database."""
-    cypher_query: str
-
-
-class GraphSchemaTool(ToolMessage):
-    request: str = "get_schema"
-    purpose: str = """To get the schema of the graph database."""
-
-
 class Neo4jSettings(BaseSettings):
     uri: str = ""
     username: str = ""
@@ -65,17 +56,22 @@ class Neo4jSettings(BaseSettings):
 
 class QueryResult(BaseModel):
     success: bool
-    data: Optional[Union[str, List[Dict[Any, Any]]]] = None
+    data: List[Dict[Any, Any]] | str | None = None
 
 
 class Neo4jChatAgentConfig(ChatAgentConfig):
     neo4j_settings: Neo4jSettings = Neo4jSettings()
     system_message: str = DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE
-    kg_schema: Optional[List[Dict[str, Any]]]
+    kg_schema: Optional[List[Dict[str, Any]]] = None
     database_created: bool = False
+    # whether agent MUST use schema_tools to get schema, i.e.
+    # schema is NOT initially provided
     use_schema_tools: bool = True
     use_functions_api: bool = True
     use_tools: bool = False
+    # whether the agent is used in a continuous chat with user,
+    # as opposed to returning a result from the task.run()
+    chat_mode: bool = False
     addressing_prefix: str = ""
 
 
@@ -89,21 +85,48 @@ class Neo4jChatAgent(ChatAgent):
         self.config: Neo4jChatAgentConfig = config
         self._validate_config()
         self._import_neo4j()
-        self._initialize_connection()
-        self._init_tool_messages()
+        self._initialize_db()
+        self._init_tools_sys_message()
         self.init_state()
 
     def init_state(self) -> None:
         super().init_state()
         self.current_retrieval_cypher_query: str = ""
+        self.tried_schema: bool = False
 
     def handle_message_fallback(
         self, msg: str | ChatDocument
-    ) -> str | ChatDocument | None:
-        """When LLM sends a no-tool msg, assume user is the intended recipient."""
+    ) -> str | ForwardTool | None:
+        """
+        When LLM sends a no-tool msg, assume user is the intended recipient,
+        and if in interactive mode, forward the msg to the user.
+        """
+
+        done_tool_name = DoneTool.default_value("request")
+        forward_tool_name = ForwardTool.default_value("request")
         if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
-            msg.metadata.recipient = Entity.USER
-            return msg
+            if self.interactive:
+                return ForwardTool(agent="User")
+            else:
+                if self.config.chat_mode:
+                    return f"""
+                    Since you did not explicitly address the User, it is not clear
+                    whether:
+                    - you intend this to be the final response to the 
+                      user's query/request, in which case you must use the 
+                      `{forward_tool_name}` to indicate this.
+                    - OR, you FORGOT to use an Appropriate TOOL,
+                      in which case you should use the available tools to
+                      make progress on the user's query/request.
+                    """
+                return f"""
+                The intent of your response is not clear:
+                - if you intended this to be the final answer to the user's query,
+                    then use the `{done_tool_name}` to indicate so,
+                    with the `content` set to the answer or result.
+                - otherwise, use one of the available tools to make progress 
+                    to arrive at the final answer.
+                """
         return None
 
     def _validate_config(self) -> None:
@@ -122,16 +145,9 @@ class Neo4jChatAgent(ChatAgent):
         try:
             import neo4j
         except ImportError:
-            raise ImportError(
-                """
-                neo4j not installed. Please install it via:
-                pip install neo4j.
-                Or when installing langroid, install it with the `neo4j` extra:
-                pip install langroid[neo4j]
-                """
-            )
+            raise LangroidImportError("neo4j", "neo4j")
 
-    def _initialize_connection(self) -> None:
+    def _initialize_db(self) -> None:
         """
         Initializes a connection to the Neo4j database using the configuration settings.
         """
@@ -144,6 +160,16 @@ class Neo4jChatAgent(ChatAgent):
                     self.config.neo4j_settings.password,
                 ),
             )
+            with self.driver.session() as session:
+                result = session.run("MATCH (n) RETURN count(n) as count")
+                count = result.single()["count"]  # type: ignore
+                self.config.database_created = count > 0
+
+            # If database has data, get schema
+            if self.config.database_created:
+                # this updates self.config.kg_schema
+                self.graph_schema_tool(None)
+
         except Exception as e:
             raise ConnectionError(f"Failed to initialize Neo4j connection: {e}")
 
@@ -226,6 +252,21 @@ class Neo4jChatAgent(ChatAgent):
             QueryResult: An object representing the outcome of the query execution.
                          It contains a success flag and an optional error message.
         """
+        # Check if query contains database/collection creation patterns
+        query_upper = query.upper()
+        is_creation_query = any(
+            [
+                "CREATE" in query_upper,
+                "MERGE" in query_upper,
+                "CREATE CONSTRAINT" in query_upper,
+                "CREATE INDEX" in query_upper,
+            ]
+        )
+
+        if is_creation_query:
+            self.config.database_created = True
+            logger.info("Detected database/collection creation query")
+
         if not self.driver:
             return QueryResult(
                 success=False, data="No database connection is established."
@@ -260,7 +301,7 @@ class Neo4jChatAgent(ChatAgent):
         else:
             print("[red]Database is not deleted!")
 
-    def retrieval_query(self, msg: CypherRetrievalTool) -> str:
+    def cypher_retrieval_tool(self, msg: CypherRetrievalTool) -> str:
         """ "
         Handle a CypherRetrievalTool message by executing a Cypher query and
         returning the result.
@@ -271,13 +312,20 @@ class Neo4jChatAgent(ChatAgent):
             str: The result of executing the cypher_query.
         """
         if not self.tried_schema:
-            return """
-            You did not yet use the `get_schema` tool to get the schema 
+            return f"""
+            You did not yet use the `{graph_schema_tool_name}` tool to get the schema 
             of the neo4j knowledge-graph db. Use that tool first before using 
-            the `retrieval_query` tool, to ensure you know all the correct
+            the `{cypher_retrieval_tool_name}` tool, to ensure you know all the correct
             node labels, relationship types, and property keys available in
             the database.
             """
+        elif not self.config.database_created:
+            return f"""
+            You have not yet created the Neo4j database. 
+            Use the `{cypher_creation_tool_name}`
+            tool to create the database first before using the 
+            `{cypher_retrieval_tool_name}` tool.
+            """
         query = msg.cypher_query
         self.current_retrieval_cypher_query = query
         logger.info(f"Executing Cypher query: {query}")
@@ -291,7 +339,7 @@ class Neo4jChatAgent(ChatAgent):
             """
         return str(response.data)
 
-    def create_query(self, msg: CypherCreationTool) -> str:
+    def cypher_creation_tool(self, msg: CypherCreationTool) -> str:
         """ "
         Handle a CypherCreationTool message by executing a Cypher query and
         returning the result.
@@ -306,6 +354,7 @@ class Neo4jChatAgent(ChatAgent):
         logger.info(f"Executing Cypher query: {query}")
         response = self.write_query(query)
         if response.success:
+            self.config.database_created = True
             return "Cypher query executed successfully"
         else:
             return str(response.data)
@@ -316,7 +365,9 @@ class Neo4jChatAgent(ChatAgent):
     # The current query works well. But we could use the queries here:
     # https://github.com/neo4j/NaLLM/blob/1af09cd117ba0777d81075c597a5081583568f9f/api/
     # src/driver/neo4j.py#L30
-    def get_schema(self, msg: GraphSchemaTool | None) -> str:
+    def graph_schema_tool(
+        self, msg: GraphSchemaTool | None
+    ) -> str | Optional[Union[str, List[Dict[Any, Any]]]]:
         """
         Retrieves the schema of a Neo4j graph database.
 
@@ -334,27 +385,42 @@ class Neo4jChatAgent(ChatAgent):
              to database connectivity or query execution.
         """
         self.tried_schema = True
+        if self.config.kg_schema is not None and len(self.config.kg_schema) > 0:
+            return self.config.kg_schema
         schema_result = self.read_query("CALL db.schema.visualization()")
         if schema_result.success:
-            # ther is a possibility that the schema is empty, which is a valid response
+            # there is a possibility that the schema is empty, which is a valid response
             # the schema.data will be: [{"nodes": [], "relationships": []}]
-            return json.dumps(schema_result.data)
+            self.config.kg_schema = schema_result.data  # type: ignore
+            return schema_result.data
         else:
             return f"Failed to retrieve schema: {schema_result.data}"
 
-    def _init_tool_messages(self) -> None:
+    def _init_tools_sys_message(self) -> None:
         """Initialize message tools used for chatting."""
         self.tried_schema = False
         message = self._format_message()
         self.config.system_message = self.config.system_message.format(mode=message)
-        if self.config.addressing_prefix != "":
+        if self.config.chat_mode:
+            self.config.addressing_prefix = self.config.addressing_prefix or SEND_TO
             self.config.system_message += ADDRESSING_INSTRUCTION.format(
                 prefix=self.config.addressing_prefix
             )
+        else:
+            self.config.system_message += DONE_INSTRUCTION
         super().__init__(self.config)
-        self.enable_message(CypherRetrievalTool)
-        self.enable_message(CypherCreationTool)
-        self.enable_message(GraphSchemaTool)
+        # Note we are enabling GraphSchemaTool regardless of whether
+        # self.config.use_schema_tools 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
+        self.enable_message(
+            [
+                GraphSchemaTool,
+                CypherRetrievalTool,
+                CypherCreationTool,
+                DoneTool,
+            ]
+        )
 
     def _format_message(self) -> str:
         if self.driver is None:
@@ -363,5 +429,5 @@ class Neo4jChatAgent(ChatAgent):
         return (
             SCHEMA_TOOLS_SYS_MSG
             if self.config.use_schema_tools
-            else DEFAULT_SYS_MSG.format(schema=self.get_schema(None))
+            else SCHEMA_PROVIDED_SYS_MSG.format(schema=self.graph_schema_tool(None))
         )