Merge pull request #25 from tigergraph/GML-1496-generate-function-str…

…uctured-out

feat(generate function): structure output and include reasoning

app/main.py

@@ -140,11 +140,11 @@ def retrieve_answer(graphname, query: NaturalLanguageQuery, credentials: Annotat
         steps = agent.question_for_agent(query.query)
         logger.debug(f"/{graphname}/query request_id={req_id_cv.get()} agent executed")
         try:
-            function_call = steps["intermediate_steps"][-1][-1].split("Function ")[1].split(" produced")[0]
-            res = steps["intermediate_steps"][-1][-1].split("the result ")[-1]
+            generate_func_output = steps["intermediate_steps"][-1][-1]
             resp.natural_language_response = steps["output"]
-            resp.query_sources = {"function_call": function_call,
-                                "result": json.loads(res)}
+            resp.query_sources = {"function_call": generate_func_output["function_call"],
+                                "result": json.loads(generate_func_output["result"]),
+                                "reasoning": generate_func_output["reasoning"]}
             resp.answered_question = True
         except Exception as e:
             resp.natural_language_response = steps["output"]

app/prompts/azure_open_ai_gpt35_turbo_instruct/generate_function.txt

@@ -10,4 +10,5 @@ First Docstring: {doc1}
 Second Docstring: {doc2}
 Third Docstring: {doc3}
 
-Given a pyTigerGraph TigerGraphConnection object named `conn`, what would be the Python function to answer the question? Only answer with the function call.
+Follow the output directions below on how to structure your response:
+{format_instructions}

app/prompts/gcp_vertexai_palm/generate_function.txt

@@ -1,7 +1,6 @@
 Use the vertex types, edge types, and their attributes and IDs below to write the pyTigerGraph function call to answer the question using a pyTigerGraph connection.
 YOU ARE INTERACTING WITH A LABELED PROPERTY GRAPH DATABSE THROUGH PYTHON FUNCTION CALLS, WHERE VERTICES AND EDGES HAVE UNIQUE TYPES AND ATTRIBUTES.
-Documentation contains helpful Python docstrings for the various functions. Use this knowledge to construct the proper function call. Choose one function to execute.
-ONLY WRITE THE FUNCTION CALL, without `conn` in front of it.
+The documentation below contains helpful Python docstrings for the various functions. Use this knowledge to construct the proper function call. Choose one function to execute.
 Vertex Types: {vertex_types}
 Vertex Attributes: {vertex_attributes}
 Vertex IDs: {vertex_ids}
@@ -11,4 +10,6 @@ Question: {question}
 First Docstring: {doc1}
 Second Docstring: {doc2}
 Third Docstring: {doc3}
-Python Call: conn.
+
+Follow the output directions below on how to structure your response:
+{format_instructions}

app/prompts/openai_gpt4/generate_function.txt

@@ -1,8 +1,6 @@
 Use the vertex types, edge types, and their attributes and IDs below to write the pyTigerGraph function call to answer the question using a pyTigerGraph connection.
-YOU ARE INTERACTING WITH A LABELED PROPERTY GRAPH DATABSE THROUGH PYTHON FUNCTION CALLS, WHERE VERTICES AND EDGES HAVE UNIQUE TYPES AND ATTRIBUTES.
 When the question asks for "How many", make sure to select a function that contains "Count" in the description/function call. Make sure never to generate a function that is not listed below.
 Documentation contains helpful Python docstrings for the various functions. Use this knowledge to construct the proper function call. Choose one function to execute.
-ONLY WRITE THE FUNCTION CALL, without `conn` in front of it.
 Vertex Types: {vertex_types}
 Vertex Attributes: {vertex_attributes}
 Vertex IDs: {vertex_ids}
@@ -12,4 +10,6 @@ Question: {question}
 First Docstring: {doc1}
 Second Docstring: {doc2}
 Third Docstring: {doc3}
-Python Call: conn.
+
+Follow the output directions below on how to structure your response:
+{format_instructions}

app/schemas/tool_io_schemas.py

@@ -14,4 +14,8 @@ class AgentOutput(BaseModel):
     function_call: str = Field(description="Function call used to generate answer")
 
 class MapAttributeToAttributeResponse(BaseModel):
-    attr_map: Dict[str, str] = Field(description="The dictionary of the form {'source_attribute': 'output_attribute'}")
+    attr_map: Dict[str, str] = Field(description="The dictionary of the form {'source_attribute': 'output_attribute'}")
+
+class GenerateFunctionResponse(BaseModel):
+    connection_func_call: str = Field(description="The function call to make to answer the question. Must start with conn.")
+    func_call_reasoning: str = Field(description="The reason why the function call was generated to answer the question.")

app/tools/generate_function.py

@@ -6,7 +6,7 @@
 from langchain.output_parsers import PydanticOutputParser
 from pyTigerGraph import TigerGraphConnection
 from langchain.pydantic_v1 import BaseModel, Field, validator
-from app.schemas import MapQuestionToSchemaResponse
+from app.schemas import MapQuestionToSchemaResponse, GenerateFunctionResponse
 from typing import List, Dict, Type, Optional, Union
 from app.embedding_utils.embedding_services import EmbeddingModel
 from app.embedding_utils.embedding_stores import EmbeddingStore
@@ -75,10 +75,6 @@ def _run(self, question: str,
                     The dictionary of edge attributes the question mentions, in the form {"edge_type": ["attr1", "attr2"]}
         """
         logger.info(f"request_id={req_id_cv.get()} ENTRY GenerateFunction._run()")
-        PROMPT = PromptTemplate(
-            template=self.prompt, input_variables=["question", "vertex_types", "edge_types", "vertex_attributes",
-                                                   "vertex_ids", "edge_attributes", "doc1", "doc2", "doc3"]
-        )
 
         if target_vertex_types == [] and target_edge_types == []:
             return "No vertex or edge types recognized. MapQuestionToSchema and then try again."
@@ -101,6 +97,15 @@ def _run(self, question: str,
 
         logger.debug_pii(f"request_id={req_id_cv.get()} retrieving documents for question={lookup_question}")
 
+        func_parser = PydanticOutputParser(pydantic_object=GenerateFunctionResponse)
+
+        PROMPT = PromptTemplate(
+            template=self.prompt,
+            input_variables=["question", "vertex_types", "edge_types", "vertex_attributes",
+                             "vertex_ids", "edge_attributes", "doc1", "doc2", "doc3"],
+            partial_variables={"format_instructions": func_parser.get_format_instructions()}
+        )
+
         docs = self.embedding_store.retrieve_similar(self.embedding_model.embed_query(lookup_question), top_k=3)
         inputs = [{"question": question, 
                     "vertex_types": target_vertex_types, #self.conn.getVertexTypes(), 
@@ -118,18 +123,21 @@ def _run(self, question: str,
         chain = LLMChain(llm=self.llm, prompt=PROMPT)
         generated = chain.apply(inputs)[0]["text"]
         logger.debug(f"request_id={req_id_cv.get()} generated function")
-
+        generated = func_parser.invoke(generated)
         try:
-            generated = validate_function_call(self.conn, generated, docs)
+            parsed_func = validate_function_call(self.conn, generated.connection_func_call, docs)
         except InvalidFunctionCallException as e:
             logger.warning(f"request_id={req_id_cv.get()} EXIT GenerateFunction._run() with exception={e}")
             return e
 
         try:
             loc = {}
-            exec("res = conn."+generated, {"conn": self.conn}, loc)
+            exec("res = conn."+parsed_func, {"conn": self.conn}, loc)
             logger.info(f"request_id={req_id_cv.get()} EXIT GenerateFunction._run()")
-            return "Function {} produced the result {}".format(generated, json.dumps(loc["res"]))
+            return {"function_call": parsed_func,
+                    "result": json.dumps(loc["res"]),
+                    "reasoning": generated.func_call_reasoning}
+            #return "Function {} produced the result {}, due to reason {}".format(generated, json.dumps(loc["res"]), generated.func_call_reasoning)
         except Exception as e:
             logger.warning(f"request_id={req_id_cv.get()} EXIT GenerateFunction._run() with exception={e}")
             raise ToolException("The function {} did not execute correctly. Please rephrase your question and try again".format(generated))