langroid 0.59.0b3__py3-none-any.whl → 0.59.1__py3-none-any.whl

langroid/agent/special/lance_tools.py-e

@@ -1,61 +0,0 @@
-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-e

@@ -1,430 +0,0 @@
-import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
-
-from pydantic_settings import BaseSettings
-from rich import print
-from rich.console import Console
-
-from pydantic import BaseModel, ConfigDict
-
-if TYPE_CHECKING:
-    import neo4j
-
-from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
-from langroid.agent.chat_document import ChatDocument
-from langroid.agent.special.neo4j.system_messages import (
-    ADDRESSING_INSTRUCTION,
-    DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE,
-    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__)
-
-console = Console()
-
-NEO4J_ERROR_MSG = "There was an error in your Cypher Query"
-
-
-# TOOLS to be used by the agent
-
-
-class Neo4jSettings(BaseSettings):
-    uri: str = ""
-    username: str = ""
-    password: str = ""
-    database: str = ""
-
-    model_config = ConfigDict(env_prefix="NEO4J_")
-
-
-class QueryResult(BaseModel):
-    success: bool
-    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]]] = 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 = ""
-
-
-class Neo4jChatAgent(ChatAgent):
-    def __init__(self, config: Neo4jChatAgentConfig):
-        """Initialize the Neo4jChatAgent.
-
-        Raises:
-            ValueError: If database information is not provided in the config.
-        """
-        self.config: Neo4jChatAgentConfig = config
-        self._validate_config()
-        self._import_neo4j()
-        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 | 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:
-            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:
-        """Validate the configuration to ensure all necessary fields are present."""
-        assert isinstance(self.config, Neo4jChatAgentConfig)
-        if (
-            self.config.neo4j_settings.username is None
-            and self.config.neo4j_settings.password is None
-            and self.config.neo4j_settings.database
-        ):
-            raise ValueError("Neo4j env information must be provided")
-
-    def _import_neo4j(self) -> None:
-        """Dynamically imports the Neo4j module and sets it as a global variable."""
-        global neo4j
-        try:
-            import neo4j
-        except ImportError:
-            raise LangroidImportError("neo4j", "neo4j")
-
-    def _initialize_db(self) -> None:
-        """
-        Initializes a connection to the Neo4j database using the configuration settings.
-        """
-        try:
-            assert isinstance(self.config, Neo4jChatAgentConfig)
-            self.driver = neo4j.GraphDatabase.driver(
-                self.config.neo4j_settings.uri,
-                auth=(
-                    self.config.neo4j_settings.username,
-                    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}")
-
-    def close(self) -> None:
-        """close the connection"""
-        if self.driver:
-            self.driver.close()
-
-    def retry_query(self, e: Exception, query: str) -> str:
-        """
-        Generate an error message for a failed Cypher query and return it.
-
-        Args:
-            e (Exception): The exception raised during the Cypher query execution.
-            query (str): The Cypher query that failed.
-
-        Returns:
-            str: The error message.
-        """
-        logger.error(f"Cypher Query failed: {query}\nException: {e}")
-
-        # Construct the error message
-        error_message_template = f"""\
-        {NEO4J_ERROR_MSG}: '{query}'
-        {str(e)}
-        Run a new query, correcting the errors.
-        """
-
-        return error_message_template
-
-    def read_query(
-        self, query: str, parameters: Optional[Dict[Any, Any]] = None
-    ) -> QueryResult:
-        """
-        Executes a given Cypher query with parameters on the Neo4j database.
-
-        Args:
-            query (str): The Cypher query string to be executed.
-            parameters (Optional[Dict[Any, Any]]): A dictionary of parameters for
-                                                    the query.
-
-        Returns:
-            QueryResult: An object representing the outcome of the query execution.
-        """
-        if not self.driver:
-            return QueryResult(
-                success=False, data="No database connection is established."
-            )
-
-        try:
-            assert isinstance(self.config, Neo4jChatAgentConfig)
-            with self.driver.session(
-                database=self.config.neo4j_settings.database
-            ) as session:
-                result = session.run(query, parameters)
-                if result.peek():
-                    records = [record.data() for record in result]
-                    return QueryResult(success=True, data=records)
-                else:
-                    return QueryResult(success=True, data=[])
-        except Exception as e:
-            logger.error(f"Failed to execute query: {query}\n{e}")
-            error_message = self.retry_query(e, query)
-            return QueryResult(success=False, data=error_message)
-        finally:
-            self.close()
-
-    def write_query(
-        self, query: str, parameters: Optional[Dict[Any, Any]] = None
-    ) -> QueryResult:
-        """
-        Executes a write transaction using a given Cypher query on the Neo4j database.
-        This method should be used for queries that modify the database.
-
-        Args:
-            query (str): The Cypher query string to be executed.
-            parameters (dict, optional): A dict of parameters for the Cypher query.
-
-        Returns:
-            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."
-            )
-
-        try:
-            assert isinstance(self.config, Neo4jChatAgentConfig)
-            with self.driver.session(
-                database=self.config.neo4j_settings.database
-            ) as session:
-                session.write_transaction(lambda tx: tx.run(query, parameters))
-                return QueryResult(success=True)
-        except Exception as e:
-            logging.warning(f"An error occurred: {e}")
-            error_message = self.retry_query(e, query)
-            return QueryResult(success=False, data=error_message)
-        finally:
-            self.close()
-
-    # TODO: test under enterprise edition because community edition doesn't allow
-    # database creation/deletion
-    def remove_database(self) -> None:
-        """Deletes all nodes and relationships from the current Neo4j database."""
-        delete_query = """
-                MATCH (n)
-                DETACH DELETE n
-            """
-        response = self.write_query(delete_query)
-
-        if response.success:
-            print("[green]Database is deleted!")
-        else:
-            print("[red]Database is not deleted!")
-
-    def cypher_retrieval_tool(self, msg: CypherRetrievalTool) -> str:
-        """ "
-        Handle a CypherRetrievalTool message by executing a Cypher query and
-        returning the result.
-        Args:
-            msg (CypherRetrievalTool): The tool-message to handle.
-
-        Returns:
-            str: The result of executing the cypher_query.
-        """
-        if not self.tried_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 `{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}")
-        response = self.read_query(query)
-        if isinstance(response.data, list) and len(response.data) == 0:
-            return """
-            No results found; check if your query used the right label names -- 
-            remember these are case sensitive, so you have to use the exact label
-            names you found in the schema. 
-            Or retry using one of the  RETRY-SUGGESTIONS in your instructions. 
-            """
-        return str(response.data)
-
-    def cypher_creation_tool(self, msg: CypherCreationTool) -> str:
-        """ "
-        Handle a CypherCreationTool message by executing a Cypher query and
-        returning the result.
-        Args:
-            msg (CypherCreationTool): The tool-message to handle.
-
-        Returns:
-            str: The result of executing the cypher_query.
-        """
-        query = msg.cypher_query
-
-        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)
-
-    # TODO: There are various ways to get the schema. The current one uses the func
-    # `read_query`, which requires post processing to identify whether the response upon
-    # the schema query is valid. Another way is to isolate this func from `read_query`.
-    # 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 graph_schema_tool(
-        self, msg: GraphSchemaTool | None
-    ) -> str | Optional[Union[str, List[Dict[Any, Any]]]]:
-        """
-        Retrieves the schema of a Neo4j graph database.
-
-        Args:
-            msg (GraphSchemaTool): An instance of GraphDatabaseSchema, typically
-            containing information or parameters needed for the database query.
-
-        Returns:
-            str: The visual representation of the database schema as a string, or a
-            message stating that the database schema is empty or not valid.
-
-        Raises:
-            This function does not explicitly raise exceptions but depends on the
-            behavior of 'self.read_query' method, which might raise exceptions related
-             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:
-            # there is a possibility that the schema is empty, which is a valid response
-            # the schema.data will be: [{"nodes": [], "relationships": []}]
-            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_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.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)
-        # 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:
-            raise ValueError("Database driver None")
-        assert isinstance(self.config, Neo4jChatAgentConfig)
-        return (
-            SCHEMA_TOOLS_SYS_MSG
-            if self.config.use_schema_tools
-            else SCHEMA_PROVIDED_SYS_MSG.format(schema=self.graph_schema_tool(None))
-        )