diff --git a/README.md b/README.md index 7abe474..cc866ed 100644 --- a/README.md +++ b/README.md @@ -219,7 +219,7 @@ In addition to the `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_API_KEY`, and `azure_d } } } - ``` +``` * AWS Bedrock @@ -276,7 +276,7 @@ In addition to the `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_API_KEY`, and `azure_d * Hugging Face -Example configuration for a model on Hugging Face with a dedicated endpoint is shown below. Please specify your configuration details:\ +Example configuration for a model on Hugging Face with a dedicated endpoint is shown below. Please specify your configuration details: ```json { diff --git a/common/db/connections.py b/common/db/connections.py index 095c54a..f8fb00d 100644 --- a/common/db/connections.py +++ b/common/db/connections.py @@ -160,3 +160,43 @@ def elevate_db_connection_to_token(host, username, password, graphname, async_co conn.restppUrl = conn.restppUrl+"/restpp" return conn + +def get_schema_ver(conn: TigerGraphConnectionProxy) -> int: + """Retrieves the schema version of the graph by running an interpreted query. + + Returns: + The schema version as an integer. + """ + logger.info("entry: _get_schema_ver") + + # Create the interpreted query to get schema version + query_text = f'INTERPRET QUERY () FOR GRAPH {conn.graphname} {{ PRINT "OK"; }}' + + try: + # Run the interpreted query + #result = self.conn.runInterpretedQuery(query_text) + if conn._version_greater_than_4_0(): + ret = conn._post(conn.gsUrl + "/gsql/v1/queries/interpret", + params={}, data=query_text, authMode="pwd", resKey="version", + headers={'Content-Type': 'text/plain'}) + else: + ret = conn._post(conn.gsUrl + "/gsqlserver/interpreted_query", data=query_text, + params={}, authMode="pwd", resKey="version") + + schema_version_int = None + if ret and "version" in ret: + version_info = ret["version"] + if isinstance(version_info, dict) and "schema" in version_info: + schema_version = version_info["schema"] + try: + schema_version_int = int(schema_version) + except (ValueError, TypeError): + logger.warning(f"Schema version '{schema_version}' could not be converted to integer") + if schema_version_int is None: + logger.warning("Schema version not found in query result") + logger.info("exit: _get_schema_ver") + return schema_version_int + + except Exception as e: + logger.error(f"Error getting schema version: {str(e)}") + raise Exception(f"Failed to get schema version: {str(e)}") diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 7911173..730774d 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -25,7 +25,7 @@ def generate_function_prompt(self): @property def generate_cypher_prompt(self): """Property to get the prompt for the GenerateCypher tool.""" - prompt = """You're an expert in OpenCypher programming. Given the following schema: {schema}, what is the OpenCypher query that retrieves the {question} + prompt = """You're an expert in OpenCypher programming. Given the following schema, what is the OpenCypher query that retrieves the {question} Only include attributes that are found in the schema. Never include any attributes that are not found in the schema. Use attributes instead of primary id if attribute name is closer to the keyword type in the question. Use as less vertex type, edge type and attributes as possible. If an attribute is not found in the schema, please exclude it from the query. @@ -33,6 +33,11 @@ def generate_cypher_prompt(self): Never use directed edge pattern in the OpenCypher query. Always use and create query using undirected pattern. Always use double quotes for strings instead of single quotes. + Avoid generating invalid OpenCypher queries based on the errors from history below. + + Schema: {schema} + History: {history} + You cannot use the following clauses: OPTIONAL MATCH CREATE @@ -48,6 +53,42 @@ def generate_cypher_prompt(self): ONLY write the OpenCypher query in the response. Do not include any other information in the response.""" return prompt + @property + def generate_gsql_prompt(self): + """Property to get the prompt for the GenerateGSQL tool.""" + prompt = """You're an expert in GSQL (Graph SQL) programming for TigerGraph. Given the following schema: {schema}, what is the GSQL query that retrieves the answer for question: {question} + Only include attributes that are found in the schema. Never include any attributes that are not found in the schema. + Use attributes instead of primary id if attribute name is more similar to the keyword type in the question. + Use as few vertex types, edge types and attributes as possible. If an attribute is not found in the schema, please exclude it from the query. + Do not return attributes that are not explicitly mentioned in the question. If a vertex type is mentioned in the question, only return the vertex. + Always use double quotes for strings instead of single quotes. + Use alias for ORDER BY if any, and make sure the alias or attributes used in ORDER BY is also in PRINT. Always add ASC or DESC for ORDER BY based on data type. + + Avoid generating invalid GSQL queries based on the errors from history below. + + Schema: {schema} + History: {history} + + Additionally, you cannot use the following clauses: + CREATE + DELETE + INSERT + UPDATE + UPSERT + + Here's some commonly used abbreviations: + dt -> date + pct -> percentage + qty -> quantity + lng -> longitude + cm -> Contract Manufacturer + + Always make the GSQL query returns the entity in the original question together with the data to be queried. + Make sure to have correct attribute names in the GSQL query and not to name result aliases that are vertex or edge types, operator or function names, and other reserved keywords, always construct alias with multiple words connected with underscore. + + ONLY write the GSQL query in the response. Do not include any other information in the response.""" + return prompt + @property def route_response_prompt(self): """Property to get the prompt for the RouteResponse tool.""" diff --git a/common/llm_services/google_genai_service.py b/common/llm_services/google_genai_service.py index a96716d..7b6c7cf 100644 --- a/common/llm_services/google_genai_service.py +++ b/common/llm_services/google_genai_service.py @@ -47,6 +47,14 @@ def generate_cypher_prompt(self): else: return super().generate_cypher_prompt + @property + def generate_gsql_prompt(self): + filepath = self.prompt_path + "generate_gsql.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().generate_gsql_prompt + @property def entity_relationship_extraction_prompt(self): return self._read_prompt_file( diff --git a/common/llm_services/openai_service.py b/common/llm_services/openai_service.py index 6fa5d75..8834353 100644 --- a/common/llm_services/openai_service.py +++ b/common/llm_services/openai_service.py @@ -46,6 +46,14 @@ def generate_cypher_prompt(self): else: return super().generate_cypher_prompt + @property + def generate_gsql_prompt(self): + filepath = self.prompt_path + "generate_gsql.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().generate_gsql_prompt + @property def entity_relationship_extraction_prompt(self): return self._read_prompt_file( diff --git a/common/prompts/openai_gpt4/chatbot_response.txt b/common/prompts/openai_gpt4/chatbot_response.txt index c412948..f0a8311 100644 --- a/common/prompts/openai_gpt4/chatbot_response.txt +++ b/common/prompts/openai_gpt4/chatbot_response.txt @@ -1,12 +1,12 @@ You are a highly efficient and empathetic AI-powered customer support agent for TigerGraph. Your goal is to provide accurate, helpful, and friendly support while maintaining professionalism. Follow these guidelines: 1. **Understand the User's Needs**: Carefully interpret the user's question or issue. Ask clarifying questions if necessary. -2. **Provide Accurate Answers**: Use the given context in JSON format to deliver reliable and precise responses. +2. **Provide Accurate Answers**: Use the given context in JSON format as well as the data structure of the original query used to fecth the context to deliver reliable and precise responses. 3. **Be Polite and Friendly**: Maintain a tone that is warm, respectful, and supportive. 4. **Stay Within Scope**: Respond only with information relevant to the user's question. If the information is unavailable, apologize and suggest alternative actions. 5. **Encourage Next Steps**: Proactively suggest solutions or next steps, such as contacting a human agent if the issue cannot be resolved by the AI. 6. **Escalate When Needed**: Recognize complex or sensitive queries that require human attention and politely guide the user to contact human support. -7. **Reference Relevant Knowledge**: Ground your responses based on the context and resources. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID]. +7. **Reference Relevant Knowledge**: Ground your responses based on the context and resources, and referencing the data struction from the query providedy. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID]. 8. **Here are the channels to reach the TigerGraph team**: - [Contact Sales](https://www.tigergraph.com/contact/) - [Support Portal](https://www.tigergraph.com/support/) @@ -29,4 +29,5 @@ Make sure all the fields required by the format instructions are included, set a Question: {question} Context: {context} +Query: {query} Format: {format_instructions} diff --git a/common/prompts/openai_gpt4/generate_cypher.txt b/common/prompts/openai_gpt4/generate_cypher.txt index 91fcbe6..a4d920c 100644 --- a/common/prompts/openai_gpt4/generate_cypher.txt +++ b/common/prompts/openai_gpt4/generate_cypher.txt @@ -1,12 +1,18 @@ -You're an expert in OpenCypher programming. Given the following schema: {schema}, what is the OpenCypher query that retrieves the answer for question: {question}. +You're an expert in OpenCypher programming. Given the following schema, find the best OpenCypher query that retrieves the answer for question {question}. If there're multiple words in the question having same meaning then remove the duplication. Only include attributes that are found in the schema. Never include any attributes that are not found in the schema. Use attributes instead of primary id if attribute name is more similar to the keyword type in the question. Use as less vertex type, edge type and attributes as possible. If an attribute is not found in the schema, please exclude it from the query. +Always make sure the attributes used exist in the vertex type or edge type referenced, DO NOT use an attribute that does not exist in the vertex or edge from the schema. Do not return attributes that are not explicitly mentioned in the question. If a vertex type is mentioned in the question, only return the vertex. Never use directed edge pattern in the OpenCypher query. Always use and create query using undirected pattern. Always ensure the edge used starts from and ends with correct vertex types matching the schema. Always use double quotes for strings instead of single quotes. -Use alias for ORDER BY if any, and make sure the same alias is used in both RETURN and ORDER BY. Always add ASC or DESC for ORDER BY based on data type. +Use alias for ORDER BY if any, and make sure the alias or attributes used in ORDER BY is also in RETURN. Always add ASC or DESC for ORDER BY based on data type. + +Avoid to generate invalid OpenCypher queries based on the errors from history below. + +Schema: {schema} +History: {history} Only use the Supported Clauses, Operators, Functions and Expressions below but do not use any of the Unsupported Features, Functions or Syntax Limitations below: @@ -33,6 +39,9 @@ Others: id(), elementId(), labels(), properties(), timestamp() Supported Expressions: CASE: Conditional logic. +Supported Operators: +Comparison: IS NULL, IS NOT NULL + Unsupported Features: Clauses Not Yet Supported CALL, CREATE, MERGE, REMOVE, SET, UNION, UNION ALL, UNWIND @@ -55,6 +64,14 @@ UNION ALL UNWIND SET -Make sure to have correct attribute names in the OpenCypher query and not to name result aliases that are vertex or edge types, operator or function names, and other reserved keywords. +Here's some commonly used abbreviations: +dt -> date +pct -> percentage +qty -> quantity +lng -> longitude +cm -> Contract Manufacturer + +Always make the cypher query returns the entity in the original question together with the data to be queried. +Make sure to have correct attribute names in the OpenCypher query and not to name result aliases that are vertex or edge types, operator or function names, and other reserved keywords, always construct alias with multiple words connected with underscore. ONLY write the OpenCypher query in the response. Do not include any other information in the response. diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py index c6d56e9..f9f86e2 100644 --- a/graphrag/app/agent/agent_generation.py +++ b/graphrag/app/agent/agent_generation.py @@ -17,11 +17,12 @@ class TigerGraphAgentGenerator: def __init__(self, llm_model): self.llm = llm_model - def generate_answer(self, question: str, context: str) -> dict: + def generate_answer(self, question: str, context: str, query: str = "") -> dict: """Generate an answer based on the question and context. Args: question: str: The question to generate an answer for. context: str: The context to generate an answer from. + query: str: The original query used to fetch the conext. Returns: str: The answer to the question. """ @@ -31,7 +32,7 @@ def generate_answer(self, question: str, context: str) -> dict: prompt = PromptTemplate( template=self.llm.chatbot_response_prompt, - input_variables=["question", "context"], + input_variables=["question", "context", "query"], partial_variables={ "format_instructions": answer_parser.get_format_instructions() } @@ -40,6 +41,7 @@ def generate_answer(self, question: str, context: str) -> dict: full_prompt = prompt.format( question=question, context=context, + query=query, format_instructions=answer_parser.get_format_instructions() ) @@ -48,7 +50,7 @@ def generate_answer(self, question: str, context: str) -> dict: usage_data = {} with get_openai_callback() as cb: - generation = rag_chain.invoke({"question": question, "context": context}) + generation = rag_chain.invoke({"question": question, "context": context, "query": query}) usage_data["input_tokens"] = cb.prompt_tokens usage_data["output_tokens"] = cb.completion_tokens diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index ae0d8e5..3d2b93b 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -163,14 +163,22 @@ def generate_cypher(self, state): Run the agent cypher generator. """ self.emit_progress("Generating the Cypher to answer your question") - cypher = self.cypher_gen._run(state["question"]) - logger.info(f"cypher: {cypher}") + gen_history = [] + response_json = None - response = self.db_connection.gsql(cypher) - response_lines = response.split("\n") - try: + for i in range(3): + cypher = self.cypher_gen._run(state["question"], gen_history) + logger.info(f"cypher: {cypher}") + + response = self.db_connection.gsql(cypher) + response_lines = response.split("\n") json_str = "\n".join(response_lines[1:]) - response_json = json.loads(json_str) + try: + response_json = json.loads(json_str) + break + except Exception as e: + gen_history.append(f"{i}: {cypher}\n\tError: {json_str}\n") + if response_json: state["context"] = { "answer": response_json["results"][0], "cypher": cypher, @@ -178,10 +186,11 @@ def generate_cypher(self, state): cypher ), } - except Exception as e: + else: state["context"] = { "error": True, "cypher": cypher, + "answer": json_str } if state["error_history"] is None: state["error_history"] = [] @@ -355,7 +364,7 @@ def generate_answer(self, state): logger.debug_pii( f"""request_id={req_id_cv.get()} Got result: {state["context"]["answer"]}""" ) - answer = step.generate_answer(state["question"], state["context"]["answer"]) + answer = step.generate_answer(state["question"], state["context"]["answer"], state["context"]["cypher"]) logger.debug_pii( f"request_id={req_id_cv.get()} Generated answer: {answer.generated_answer}" ) diff --git a/graphrag/app/tools/__init__.py b/graphrag/app/tools/__init__.py index 4f42ec8..76c0fc1 100644 --- a/graphrag/app/tools/__init__.py +++ b/graphrag/app/tools/__init__.py @@ -1,5 +1,6 @@ from .generate_function import GenerateFunction from .map_question_to_schema import MapQuestionToSchema from .generate_cypher import GenerateCypher +from .generate_gsql import GenerateGSQL from .validation_utils import MapQuestionToSchemaException, InvalidFunctionCallException from .validation_utils import validate_schema, validate_function_call diff --git a/graphrag/app/tools/generate_cypher.py b/graphrag/app/tools/generate_cypher.py index 009f746..cfef440 100644 --- a/graphrag/app/tools/generate_cypher.py +++ b/graphrag/app/tools/generate_cypher.py @@ -1,10 +1,12 @@ import logging +from typing import Iterable from langchain_community.callbacks.manager import get_openai_callback from langchain_core.output_parsers import StrOutputParser from langchain.prompts import PromptTemplate from langchain.tools import BaseTool from langchain.llms.base import LLM from common.metrics.tg_proxy import TigerGraphConnectionProxy +from common.db.connections import get_schema_ver logger = logging.getLogger(__name__) @@ -17,6 +19,8 @@ class GenerateCypher(BaseTool): description: str = "Generates a Cypher query for the question." conn: TigerGraphConnectionProxy = None llm: LLM = None + schema_rep: str = None + schema_ver: int = None def __init__(self, conn: TigerGraphConnectionProxy, llm): """Initialize GenerateCypher. @@ -31,8 +35,14 @@ def __init__(self, conn: TigerGraphConnectionProxy, llm): super().__init__() self.conn = conn self.llm = llm + self.schema_rep = "" + self.schema_ver = -1 def _generate_schema_rep(self): + schema_ver = get_schema_ver(self.conn) + if schema_ver is not None and self.schema_ver == schema_ver: + logger.info(f"Reusing existing schema rep for schema version {schema_ver}") + return self.schema_rep verts = self.conn.getVertexTypes() edges = self.conn.getEdgeTypes() vertex_schema = [] @@ -63,16 +73,17 @@ def _generate_schema_rep(self): edge_info = f"""From Vertex: {from_vertex}\n\tTo Vertex: {to_vertex}""" edge_schema.append(f"""{edge}\n\t{edge_info}\n\tEdge direction: {direction}\n\tAttributes: \n\t\t{attributes}""") - schema_rep = f"""The schema of the graph is as follows: + self.schema_rep = f"""The schema of the graph is as follows: Vertex Types: {chr(10).join(vertex_schema)} Edge Types: {chr(10).join(edge_schema)} """ - return schema_rep + self.schema_ver = schema_ver if schema_ver is not None else -1 + return self.schema_rep - def generate_cypher(self, question: str) -> str: + def generate_cypher(self, question: str, history: Iterable[str]) -> str: """Generate Cypher query for the question. Args: question (str): @@ -85,18 +96,19 @@ def generate_cypher(self, question: str) -> str: template=self.llm.generate_cypher_prompt, input_variables=[ "question", - "schema" + "schema", + "history" ] ) schema = self._generate_schema_rep() - logger.debug_pii("Prompt to LLM:\n" + PROMPT.invoke({"question": question, "schema": schema}).to_string()) + logger.debug_pii("Prompt to LLM:\n" + PROMPT.invoke({"question": question, "schema": schema, "history": history}).to_string()) chain = PROMPT | self.llm.model | StrOutputParser() usage_data = {} with get_openai_callback() as cb: - out = chain.invoke({"question": question, "schema": schema}).strip("```cypher").strip("```") + out = chain.invoke({"question": question, "schema": schema, "history": history}).strip("```cypher").strip("```") usage_data["input_tokens"] = cb.prompt_tokens usage_data["output_tokens"] = cb.completion_tokens @@ -108,7 +120,7 @@ def generate_cypher(self, question: str) -> str: query_footer = "\n}" return query_header + out + query_footer - def _run(self, question: str): + def _run(self, question: str, history: Iterable[str]): """Run the GenerateCypher tool. Args: question (str): @@ -117,7 +129,7 @@ def _run(self, question: str): str: Cypher query for the question. """ - return self.generate_cypher(question) + return self.generate_cypher(question, history) - def _arun(self, question: str): + def _arun(self, question: str, history: Iterable[str]): raise NotImplementedError("Asynchronous execution is not supported for this tool.") diff --git a/graphrag/app/tools/generate_gsql.py b/graphrag/app/tools/generate_gsql.py new file mode 100644 index 0000000..3bcbe98 --- /dev/null +++ b/graphrag/app/tools/generate_gsql.py @@ -0,0 +1,139 @@ +import logging +from typing import Iterable +from langchain_community.callbacks.manager import get_openai_callback +from langchain_core.output_parsers import StrOutputParser +from langchain.prompts import PromptTemplate +from langchain.tools import BaseTool +from langchain.llms.base import LLM +from common.metrics.tg_proxy import TigerGraphConnectionProxy +from common.db.connections import get_schema_ver + +logger = logging.getLogger(__name__) + + +class GenerateGSQL(BaseTool): + """GenerateGSQL Tool. + Tool to generate and execute the appropriate GSQL query for the question. + """ + name: str = "GenerateGSQL" + description: str = "Generates a GSQL query for the question." + conn: TigerGraphConnectionProxy = None + llm: LLM = None + schema_rep: str = None + schema_ver: int = 0 + + def __init__(self, conn: TigerGraphConnectionProxy, llm): + """Initialize GenerateGSQL. + Args: + conn (TigerGraphConnection): + pyTigerGraph TigerGraphConnection connection to the appropriate database/graph with correct permissions + llm (LLM_Model): + LLM_Model class to interact with an external LLM API. + prompt (str): + prompt to use with the LLM_Model. Varies depending on LLM service. + """ + super().__init__() + self.conn = conn + self.llm = llm + self.schema_rep = "" + self.schema_ver = 0 + + def _generate_schema_rep(self): + schema_ver = get_schema_ver(self.conn) + if self.schema_rep and self.schema_ver == schema_ver: + logger.info(f"Reusing existing schema rep for schema version {schema_ver}") + return self.schema_rep + verts = self.conn.getVertexTypes() + edges = self.conn.getEdgeTypes() + vertex_schema = [] + for vert in verts: + primary_id = self.conn.getVertexType(vert)["PrimaryId"]["AttributeName"] + attributes = "\n\t\t".join([attr["AttributeName"] + " of type " + attr["AttributeType"]["Name"] + for attr in self.conn.getVertexType(vert)["Attributes"]]) + if attributes == "": + attributes = "No attributes" + vertex_schema.append(f"{vert}\n\tPrimary Id Attribute: {primary_id}\n\tAttributes: \n\t\t{attributes}") + + edge_schema = [] + for edge in edges: + from_vertex = self.conn.getEdgeType(edge)["FromVertexTypeName"] + to_vertex = self.conn.getEdgeType(edge)["ToVertexTypeName"] + direction = "Directed" if self.conn.getEdgeType(edge)["IsDirected"] else "Undirected" + #reverse_edge = conn.getEdgeType(edge)["Config"].get("REVERSE_EDGE") + attributes = "\n\t\t".join([attr["AttributeName"] + " of type " + attr["AttributeType"]["Name"] + for attr in self.conn.getEdgeType(edge)["Attributes"]]) + if attributes == "": + attributes = "No attributes" + if from_vertex == "*" or to_vertex == "*": + edge_pairs = self.conn.getEdgeType(edge)["EdgePairs"] + for an_edge in edge_pairs: + edge_info = f"""From Vertex: {an_edge["From"]}\n\tTo Vertex: {an_edge["To"]}""" + edge_schema.append(f"""{edge}\n\t{edge_info}\n\tEdge direction: {direction}\n\tAttributes: \n\t\t{attributes}""") + else: + edge_info = f"""From Vertex: {from_vertex}\n\tTo Vertex: {to_vertex}""" + edge_schema.append(f"""{edge}\n\t{edge_info}\n\tEdge direction: {direction}\n\tAttributes: \n\t\t{attributes}""") + + self.schema_rep = f"""The schema of the graph is as follows: +Vertex Types: +{chr(10).join(vertex_schema)} + +Edge Types: +{chr(10).join(edge_schema)} +""" + self.schema_ver = schema_ver + return self.schema_rep + + def generate_gsql(self, question: str, history: Iterable[str]) -> str: + """Generate GSQL query for the question. + Args: + question (str): + question to generate the GSQL query for. + history (Iterable[str]): + conversation history for context. + Returns: + str: + GSQL query for the question. + """ + PROMPT = PromptTemplate( + template=self.llm.generate_gsql_prompt, + input_variables=[ + "question", + "schema", + "history" + ] + ) + + schema = self._generate_schema_rep() + + logger.debug_pii("Prompt to LLM:\n" + PROMPT.invoke({"question": question, "schema": schema, "history": history}).to_string()) + + chain = PROMPT | self.llm.model | StrOutputParser() + usage_data = {} + with get_openai_callback() as cb: + out = chain.invoke({"question": question, "schema": schema, "history": history}).strip("```gsql").strip("```") + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"generate_gsql usage: {usage_data}") + + query_header = "USE GRAPH " + self.conn.graphname + " "+ "\n" + "INTERPRET QUERY () FOR GRAPH " + self.conn.graphname + " {" + "\n" + query_footer = "\n}" + return query_header + out + query_footer + + def _run(self, question: str, history: Iterable[str]): + """Run the GenerateGSQL tool. + Args: + question (str): + question to generate the GSQL query for. + history (Iterable[str]): + conversation history for context. + Returns: + str: + GSQL query for the question. + """ + return self.generate_gsql(question, history) + + def _arun(self, question: str, history: Iterable[str]): + raise NotImplementedError("Asynchronous execution is not supported for this tool.") \ No newline at end of file diff --git a/graphrag/app/tools/map_question_to_schema.py b/graphrag/app/tools/map_question_to_schema.py index 13c62a7..e533864 100644 --- a/graphrag/app/tools/map_question_to_schema.py +++ b/graphrag/app/tools/map_question_to_schema.py @@ -13,6 +13,7 @@ import logging from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter +from common.db.connections import get_schema_ver logger = logging.getLogger(__name__) @@ -28,6 +29,10 @@ class MapQuestionToSchema(BaseTool): llm: LLM = None prompt: str = None handle_tool_error: bool = True + schema_ver: int = None + vertices_info: list[dict] = None + edges_info: list[dict] = None + def __init__(self, conn, llm): """Initialize MapQuestionToSchema. @@ -43,6 +48,9 @@ def __init__(self, conn, llm): logger.debug(f"request_id={req_id_cv.get()} MapQuestionToSchema instantiated") self.conn = conn self.llm = llm + self.schema_ver = -1 + self.vertices_info = [] + self.edges_info = [] def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: @@ -69,31 +77,32 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: restate_chain = RESTATE_QUESTION_PROMPT | self.llm.model | parser - vertices = self.conn.getVertexTypes() - edges = self.conn.getEdgeTypes() + schema_ver = get_schema_ver(self.conn) + if schema_ver is None or self.schema_ver != schema_ver: + self.schema_ver = schema_ver if schema_ver is not None else -1 + vertices = self.conn.getVertexTypes() + edges = self.conn.getEdgeTypes() - vertices_info = [] - for vertex in vertices: - vertex_attrs = self.conn.getVertexAttrs(vertex) - attributes = [attr[0] for attr in vertex_attrs] - vertex_info = {"vertex": vertex, "attributes": attributes} - vertices_info.append(vertex_info) + for vertex in vertices: + vertex_attrs = self.conn.getVertexAttrs(vertex) + attributes = [attr[0] for attr in vertex_attrs] + vertex_info = {"vertex": vertex, "attributes": attributes} + self.vertices_info.append(vertex_info) - edges_info = [] - for edge in edges: - source_vertex = self.conn.getEdgeSourceVertexType(edge) - target_vertex = self.conn.getEdgeTargetVertexType(edge) - edge_info = {"edge": edge, "source": source_vertex, "target": target_vertex} - edges_info.append(edge_info) + for edge in edges: + source_vertex = self.conn.getEdgeSourceVertexType(edge) + target_vertex = self.conn.getEdgeTargetVertexType(edge) + edge_info = {"edge": edge, "source": source_vertex, "target": target_vertex} + self.edges_info.append(edge_info) usage_data = {} with get_openai_callback() as cb: parsed_q = restate_chain.invoke( { "vertices": vertices, - "verticesAttrs": vertices_info, + "verticesAttrs": self.vertices_info, "edges": edges, - "edgesInfo": edges_info, + "edgesInfo": self.edges_info, "question": query, "conversation": conversation, } @@ -138,9 +147,10 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: usage_data["output_tokens"] += cb.completion_tokens usage_data["total_tokens"] += cb.total_tokens usage_data["cost"] += cb.total_cost - parsed_q.target_vertex_attributes[vertex] = [ - parsed_map.get(x) for x in list(parsed_q.target_vertex_attributes[vertex]) - ] + if parsed_map: + parsed_q.target_vertex_attributes[vertex] = [ + parsed_map.get(x) for x in list(parsed_q.target_vertex_attributes[vertex]) + ] logger.debug(f"request_id={req_id_cv.get()} MapVertexAttributes applied") @@ -157,9 +167,10 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: usage_data["output_tokens"] += cb.completion_tokens usage_data["total_tokens"] += cb.total_tokens usage_data["cost"] += cb.total_cost - parsed_q.target_edge_attributes[edge] = [ - parsed_map[x] for x in list(parsed_q.target_edge_attributes[edge]) - ] + if parsed_map: + parsed_q.target_edge_attributes[edge] = [ + parsed_map[x] for x in list(parsed_q.target_edge_attributes[edge]) + ] logger.debug(f"request_id={req_id_cv.get()} MapEdgeAttributes applied")