langroid 0.1.139__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/agent/special/lance_tools.py

@@ -0,0 +1,44 @@
+import logging
+
+from pydantic import BaseModel
+
+from langroid.agent.tool_message import ToolMessage
+
+logger = logging.getLogger(__name__)
+
+
+class QueryPlan(BaseModel):
+    original_query: str
+    query: str
+    filter: str
+    dataframe_calc: str = ""
+
+
+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 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 QueryPlanAnswerTool(ToolMessage):
+    request = "query_plan_answer"  # the agent method name that handles this tool
+    purpose = """
+    Assemble query <plan> and <answer>
+    """
+    plan: QueryPlan
+    answer: str
+
+
+class QueryPlanFeedbackTool(ToolMessage):
+    request = "query_plan_feedback"
+    purpose = "To give <feedback> regarding the query plan."
+    feedback: str

langroid/agent/special/neo4j/csv_kg_chat.py

@@ -0,0 +1,174 @@
+from typing import List, Optional
+
+import pandas as pd
+import typer
+
+from langroid.agent.special.neo4j.neo4j_chat_agent import (
+    Neo4jChatAgent,
+    Neo4jChatAgentConfig,
+)
+from langroid.agent.tool_message import ToolMessage
+from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
+from langroid.parsing.table_loader import read_tabular_data
+from langroid.utils.output import status
+from langroid.vector_store.base import VectorStoreConfig
+
+app = typer.Typer()
+
+
+BUILD_KG_INSTRUCTIONS = """
+    Your task is to build a knowledge graph based on a CSV file. 
+    
+    You need to generate the graph database based on this
+    header: 
+    
+    {header}
+    
+    and these sample rows: 
+    
+    {sample_rows}. 
+    
+    Leverage the above information to: 
+    - Define node labels and their properties
+    - Infer relationships
+    - Infer constraints 
+    ASK me if you need further information to figure out the schema.
+    You can use the tool/function `pandas_to_kg` to display and confirm 
+    the nodes and relationships.
+"""
+
+DEFAULT_CSV_KG_CHAT_SYSTEM_MESSAGE = """
+    You are an expert in Knowledge Graphs and analyzing them using Neo4j.
+    You will be asked to answer questions based on the knowledge graph.
+"""
+
+
+def _preprocess_dataframe_for_neo4j(
+    df: pd.DataFrame, default_value: Optional[str] = None, remove_null_rows: bool = True
+) -> pd.DataFrame:
+    """
+    Preprocess a DataFrame for Neo4j import by fixing mismatched quotes in string
+        columns and handling null or missing values.
+
+    Args:
+        df (DataFrame): The DataFrame to be preprocessed.
+        default_value (str, optional): The default value to replace null values.
+        This is ignored if remove_null_rows is True. Defaults to None.
+        remove_null_rows (bool, optional): If True, rows with any null values will
+            be removed.
+        If False, null values will be filled with default_value. Defaults to False.
+
+    Returns:
+        DataFrame: The preprocessed DataFrame ready for Neo4j import.
+    """
+
+    # Fix mismatched quotes in string columns
+    for column in df.select_dtypes(include=["object"]):
+        df[column] = df[column].apply(
+            lambda x: x + '"' if (isinstance(x, str) and x.count('"') % 2 != 0) else x
+        )
+
+    # Handle null or missing values
+    if remove_null_rows:
+        df = df.dropna()
+    else:
+        if default_value is not None:
+            df = df.fillna(default_value)
+
+    return df
+
+
+class CSVGraphAgentConfig(Neo4jChatAgentConfig):
+    system_message: str = DEFAULT_CSV_KG_CHAT_SYSTEM_MESSAGE
+    data: str | pd.DataFrame | None  # data file, URL, or DataFrame
+    separator: None | str = None  # separator for data file
+    vecdb: None | VectorStoreConfig = None
+    llm: OpenAIGPTConfig = OpenAIGPTConfig(
+        chat_model=OpenAIChatModel.GPT4_TURBO,
+    )
+
+
+class PandasToKGTool(ToolMessage):
+    request: str = "pandas_to_kg"
+    purpose: str = """Use this tool to create ONLY nodes and their relationships based
+    on the created model.
+    Take into account that the Cypher query will be executed while iterating 
+    over the rows in the CSV file (e.g. `index, row in df.iterrows()`),
+    so there NO NEED to load the CSV.
+    Make sure you send me the cypher query in this format: 
+    - placeholders in <cypherQuery> should be based on the CSV header. 
+    - <args> an array wherein each element corresponds to a placeholder in the 
+    <cypherQuery> and provided in the same order as the headers. 
+    SO the <args> should be the result of: `[row_dict[header] for header in headers]`
+    """
+    cypherQuery: str
+    args: list[str]
+
+    @classmethod
+    def examples(cls) -> List["ToolMessage"]:
+        return [
+            cls(
+                cypherQuery="""MERGE (employee:Employee {name: $employeeName, 
+                id: $employeeId})\n
+                MERGE (department:Department {name: $departmentName})\n
+                MERGE (employee)-[:WORKS_IN]->(department)\n
+                SET employee.email = $employeeEmail""",
+                args=["employeeName", "employeeId", "departmentName", "employeeEmail"],
+            ),
+        ]
+
+
+class CSVGraphAgent(Neo4jChatAgent):
+    def __init__(self, config: CSVGraphAgentConfig):
+        formatted_build_instr = ""
+        if isinstance(config.data, pd.DataFrame):
+            df = config.data
+            self.df = df
+        else:
+            if config.data:
+                df = read_tabular_data(config.data, config.separator)
+                df_cleaned = _preprocess_dataframe_for_neo4j(df)
+
+                df_cleaned.columns = df_cleaned.columns.str.strip().str.replace(
+                    " +", "_", regex=True
+                )
+
+                self.df = df_cleaned
+
+                formatted_build_instr = BUILD_KG_INSTRUCTIONS.format(
+                    header=self.df.columns, sample_rows=self.df.head(3)
+                )
+
+        config.system_message = config.system_message + formatted_build_instr
+        super().__init__(config)
+
+        self.config: Neo4jChatAgentConfig = config
+
+        self.enable_message(PandasToKGTool)
+
+    def pandas_to_kg(self, msg: PandasToKGTool) -> str:
+        """
+        Creates nodes and relationships in the graph database based on the data in
+        a CSV file.
+
+        Args:
+            msg (PandasToKGTool): An instance of the PandasToKGTool class containing
+                the necessary information for generating nodes.
+
+        Returns:
+            str: A string indicating the success or failure of the operation.
+        """
+        with status("[cyan]Generating graph database..."):
+            if self.df is not None and hasattr(self.df, "iterrows"):
+                for counter, (index, row) in enumerate(self.df.iterrows()):
+                    row_dict = row.to_dict()
+                    response = self.write_query(
+                        msg.cypherQuery,
+                        parameters={header: row_dict[header] for header in msg.args},
+                    )
+                    # there is a possibility the generated cypher query is not correct
+                    # so we need to check the response before continuing to the
+                    # iteration
+                    if counter == 0 and not response.success:
+                        return str(response.data)
+            return "Graph database successfully generated"

langroid/agent/special/neo4j/neo4j_chat_agent.py

@@ -0,0 +1,370 @@
+import json
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, BaseSettings
+from rich import print
+from rich.console import Console
+
+from langroid.agent import ToolMessage
+
+if TYPE_CHECKING:
+    import neo4j
+
+
+from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
+from langroid.agent.chat_document import ChatDocMetaData, ChatDocument
+from langroid.agent.special.neo4j.utils.system_message import (
+    DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE,
+    DEFAULT_SYS_MSG,
+    SCHEMA_TOOLS_SYS_MSG,
+)
+from langroid.mytypes import Entity
+
+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 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 = ""
+    password: str = ""
+    database: str = ""
+
+    class Config:
+        # This enables the use of environment variables to set the settings,
+        # e.g. NEO4J_URI, NEO4J_USERNAME, etc.,
+        # which can either be set in a .env file or in the shell via export cmds.
+        env_prefix = "NEO4J_"
+
+
+class QueryResult(BaseModel):
+    success: bool
+    data: Optional[Union[str, List[Dict[Any, Any]]]] = None
+
+
+class Neo4jChatAgentConfig(ChatAgentConfig):
+    neo4j_settings: Neo4jSettings = Neo4jSettings()
+    system_message: str = DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE
+    kg_schema: Optional[List[Dict[str, Any]]]
+    database_created: bool = False
+    use_schema_tools: bool = True
+    use_functions_api: bool = True
+    use_tools: bool = False
+
+
+class Neo4jChatAgent(ChatAgent):
+    def __init__(self, config: Neo4jChatAgentConfig):
+        """Initialize the Neo4jChatAgent.
+
+        Raises:
+            ValueError: If database information is not provided in the config.
+        """
+        self.config = config
+        self._validate_config()
+        self._import_neo4j()
+        self._initialize_connection()
+        self._init_tool_messages()
+
+    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 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]
+                """
+            )
+
+    def _initialize_connection(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,
+                ),
+            )
+        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.
+        """
+        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 retrieval_query(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.
+        """
+        query = msg.cypher_query
+
+        logger.info(f"Executing Cypher query: {query}")
+        response = self.read_query(query)
+        if response.success:
+            return json.dumps(response.data)
+        else:
+            return str(response.data)
+
+    def create_query(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:
+            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 get_schema(self, msg: GraphSchemaTool | None) -> str:
+        """
+        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.
+        """
+        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
+            # the schema.data will be: [{"nodes": [], "relationships": []}]
+            return json.dumps(schema_result.data)
+        else:
+            return f"Failed to retrieve schema: {schema_result.data}"
+
+    def _init_tool_messages(self) -> None:
+        """Initialize message tools used for chatting."""
+        message = self._format_message()
+        self.config.system_message = self.config.system_message.format(mode=message)
+        super().__init__(self.config)
+        self.enable_message(CypherRetrievalTool)
+        self.enable_message(CypherCreationTool)
+        self.enable_message(GraphSchemaTool)
+
+    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 DEFAULT_SYS_MSG.format(schema=self.get_schema(None))
+        )
+
+    def agent_response(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        if msg is None:
+            return None
+
+        results = self.handle_message(msg)
+        if results is None:
+            return None
+
+        output = results
+        if NEO4J_ERROR_MSG in output:
+            output = "There was an error in the Cypher Query. Press enter to retry."
+
+        console.print(f"[red]{self.indent}", end="")
+        print(f"[red]Agent: {output}")
+        sender_name = self.config.name
+        if isinstance(msg, ChatDocument) and msg.function_call is not None:
+            sender_name = msg.function_call.name
+
+        content = results.content if isinstance(results, ChatDocument) else results
+        recipient = (
+            results.metadata.recipient if isinstance(results, ChatDocument) else ""
+        )
+
+        return ChatDocument(
+            content=content,
+            metadata=ChatDocMetaData(
+                # source=Entity.AGENT,
+                sender=Entity.LLM,
+                sender_name=sender_name,
+                recipient=recipient,
+            ),
+        )

langroid/agent/special/neo4j/utils/system_message.py

@@ -0,0 +1,46 @@
+DEFAULT_SYS_MSG = """You are a data scientist and expert in Knowledge Graphs, 
+with expertise in answering questions by interacting with a Neo4j graph database.
+
+The schema maps the Neo4j database structure. node labels, relationship types,
+and property keys available in your Neo4j database. 
+{schema}
+Do not make assumptions about the database schema before using the tools.
+Use the tool/function to learn more about the database schema."""
+
+SCHEMA_TOOLS_SYS_MSG = """You are a data scientist and expert in Knowledge Graphs, 
+with expertise in answering questions by querying Neo4j database.
+You have access to the following tools: 
+ - `retrieval_query` tool/function-call to retreive infomration from the graph database
+to answer questions.
+
+ - `create_query` tool/function-call to execute cypher query that creates
+   entities/relationships in the graph database.
+
+ - `get_schema` tool/function-call to get all the node labels, relationship 
+ types, and property keys available in your Neo4j database.
+
+You must be smart about using the right node labels, relationship types, and property
+keys based on the english description. If you are thinking of using a node label,
+relationship type, or property key that does not exist, you are probably on the wrong 
+track, so you should try your best to answer based on an existing table or column.
+DO NOT assume any nodes or relationships other than those above."""
+
+DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE = """
+{mode}
+
+You do not need to attempt answering a question with just one query. 
+You could make a sequence of Neo4j queries to help you write the final query.
+Also if you receive a null or other unexpected result,
+(a) make sure you use the available TOOLs correctly, and 
+(b) see if you have made an assumption in your Neo4j query, and try another way, 
+   or use `retrieval_query` to explore the database contents before submitting your 
+   final query. 
+(c) USE `create_query` tool/function-call to execute cypher query that creates
+   entities/relationships in the graph database.
+
+(d) USE `get_schema` tool/function-call to get all the node labels, relationship 
+ types, and property keys available in your Neo4j database.
+
+Start by asking what I would like to know about the data.
+
+"""