add query dispacher using tools calling (#309)

backend/app/rag/knowledge_graph/intent.py

@@ -3,6 +3,7 @@
 from dspy.functional import TypedChainOfThought, TypedPredictor
 from typing import List, Optional
 from pydantic import BaseModel, Field
+from llama_index.core.tools import FunctionTool
 
 from app.rag.knowledge_graph.schema import Relationship
 
@@ -13,9 +14,7 @@ class RelationshipReasoning(Relationship):
     """Relationship between two entities extracted from the query"""
 
     reasoning: str = Field(
-        description=(
-            "Explanation of the user's intention for this step."
-        )
+        description=("Explanation of the user's intention for this step.")
     )
 
 
@@ -28,23 +27,23 @@ class DecomposedFactors(BaseModel):
 
 
 class DecomposeQuery(dspy.Signature):
-    """You are a knowledge base graph expert and are very good at building knowledge graphs. Now you are assigned to extract the user's step-by-step intentions from the query.
+    """You are a knowledge base graph expert and are very good at building knowledge graphs. Now you are assigned to extract the user's step-by-step questions from the query.
 
     ## Instructions:
 
-    - Break down the user's query into a sequence of prerequisite questions (e.g., identifying specific versions) and step-by-step intentions.
-    - Represent each prerequisite and intention as a relationship: (Source Entity) - [Relationship] -> (Target Entity).
+    - Break down the user's query into a sequence of prerequisite questions (e.g., identifying specific versions) and step-by-step questions.
+    - Represent each questions as a relationship: (Source Entity) - [Relationship] -> (Target Entity).
     - Provide reasoning for each relationship, explaining the user's intention at that step.
     - Limit to no more than 5 relationships.
     - Ensure that the extracted relationships accurately reflect the user's real intentions.
     - Ensure that the relationships and intentions are grounded and factual, based on the information provided in the query.
     """
 
     query: str = dspy.InputField(
-        desc="The query text to extract the user's step-by-step intentions."
+        desc="The query text to extract the user's step-by-step questions."
     )
     factors: DecomposedFactors = dspy.OutputField(
-        desc="Representation of the user's step-by-step intentions extracted from the query."
+        desc="Representation of the user's step-by-step questions extracted from the query."
     )
 
 
@@ -67,3 +66,24 @@ def __init__(self, dspy_lm: dspy.LM, complied_program_path: Optional[str] = None
 
     def analyze(self, query: str) -> DecomposedFactors:
         return self.intent_anlysis_prog(query=query).factors
+
+    def as_tool_call(self):
+        def breakDownQuery(query: str) -> list[str]:
+            """
+            Break down the complex user query into sub-questions and solve them step-by-step.
+
+            If the query is complex, this function decomposes it into smaller sub-questions
+            and solves them individually. For simple and straightforward queries,
+            avoid calling this function.
+
+            Args:
+            query (str): The user's input query.
+
+            Returns:
+            str: The final result after solving all sub-questions.
+            """
+
+            return [query]
+
+        return FunctionTool.from_defaults(fn=breakDownQuery)
+

backend/app/rag/query_dispatcher.py

@@ -0,0 +1,44 @@
+import logging
+from typing import Optional, Sequence
+from llama_index.llms.openai import OpenAI
+from llama_index.core.tools.types import BaseTool
+from llama_index.core.tools import FunctionTool
+
+logger = logging.getLogger(__name__)
+
+DefaultSystemPrompt = """
+You are a highly skilled customer assistant, responsible for dispatching user questions to the most appropriate tools or resources. Your primary objective is to ensure each user question is handled accurately and efficiently by selecting the best-suited tool for the task.
+For more complex questions, you should break them down into clear, manageable sub-questions and route each to the relevant tools for individual resolution. It's important to maintain clarity and precision in this process, ensuring that the sub-questions are well-defined and can be resolved independently.
+If you encounter concepts or entities you are not familiar with, you can break the query down into a sub-question to clarify the specific concept or entity. For example, if the query involves “what is the latest version,” you can treat this as a sub-question to better understand the context before proceeding with the solution.
+"""
+
+class QueryDispatcher:
+    def __init__(self, llm: OpenAI, system_prompt: Optional[str] = None):
+        if system_prompt is None:
+            system_prompt = DefaultSystemPrompt
+
+        self._llm = llm
+        self._llm.system_prompt = system_prompt
+
+    def route(self, query: str, tools: Sequence["BaseTool"]) -> str:
+        response = self._llm.chat_with_tools(tools, query, allow_parallel_tool_calls=True, verbose=True)
+
+        try:
+            tool_calls = self._llm.get_tool_calls_from_response(
+                response, error_on_no_tool_call=True
+            )
+        except Exception as e:
+            logger.exception(e)
+            return f"An error occurred while processing the query: {query}"
+
+        return tool_calls
+
+
+# mock the answer process
+def answer(query: str) -> str:
+    """
+    Answer a user query. The query should be simple and straightforward.
+    """
+    return f"I need some time to answer your question: {query}."
+
+answer_tool = FunctionTool.from_defaults(fn=answer)

backend/dspy_compiled_program/decompose_query_program

@@ -37,7 +37,7 @@
         "factors": "{\"relationships\":[{\"source_entity\":\"Chat2query\",\"target_entity\":\"Error Message\",\"relationship_desc\":\"Chat2query is returning an error message saying 'Query timeout expired'.\",\"reasoning\":\"The main problem the user is facing.\"},{\"source_entity\":\"SQL Query\",\"target_entity\":\"Slow Query Log\",\"relationship_desc\":\"The reason why not to locate the SQL query in the slow query log.\",\"reasoning\":\"The secondary problem the user is facing.\"}]}"
       }
     ],
-    "signature_instructions": "You are a knowledge base graph expert and are very good at building knowledge graphs. Now you are assigned to extract the user's step-by-step intentions from the query.\n\n## Instructions:\n\n- Break down the user's query into a sequence of prerequisite questions (e.g., identifying specific versions) and step-by-step intentions.\n- Represent each prerequisite and intention as a relationship: (Source Entity) - [Relationship] -> (Target Entity).\n- Provide reasoning for each relationship, explaining the user's intention at that step.\n- Limit to no more than 5 relationships.\n- Ensure that the extracted relationships accurately reflect the user's real intentions.\n- Ensure that the relationships and intentions are grounded and factual, based on the information provided in the query.",
+    "signature_instructions": "You are a knowledge base graph expert and are very good at building knowledge graphs. Now you are assigned to extract the user's step-by-step questions from the query.\n\n## Instructions:\n\n- Break down the user's query into a sequence of prerequisite questions (e.g., identifying specific versions) and step-by-step questions.\n- Represent each questions as a relationship: (Source Entity) - [Relationship] -> (Target Entity).\n- Provide reasoning for each relationship, explaining the user's intention at that step.\n- Limit to no more than 5 relationships.\n- Ensure that the extracted relationships accurately reflect the user's real intentions.\n- Ensure that the relationships and intentions are grounded and factual, based on the information provided in the query.",
     "signature_prefix": "Factors:"
   }
 }