From 8b6097e19e5808e1a1514cdec0246a8d33f744a8 Mon Sep 17 00:00:00 2001
From: 6drf21e <nir311@163.com>
Date: Mon, 25 Mar 2024 20:35:12 +0800
Subject: [PATCH] Update get_related_questions use function calling

---
 search4all.py | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

diff --git a/search4all.py b/search4all.py
index ad6cdf0..3d0bf0c 100644
--- a/search4all.py
+++ b/search4all.py
@@ -423,7 +423,7 @@ async def server_init(_app, loop):
     )
 
 
-async def get_related_questions(_app, query, contexts):
+async def get_related_questions_bak(_app, query, contexts):
     """
     Gets related questions based on the query using the gpt-3.5-turbo model.
     """
@@ -456,6 +456,67 @@ async def get_related_questions(_app, query, contexts):
         )
         return []
 
+async def get_related_questions(_app, query, contexts):
+    """
+    Gets related questions based on the query and context.
+    """
+    _lepton_more_questions_prompt = """
+    You are a helpful assistant that helps the user to ask related questions, based on user's original question and the related contexts. Please identify worthwhile topics that can be follow-ups, and write questions no longer than 20 words each. Please make sure that specifics, like events, names, locations, are included in follow up questions so they can be asked standalone. For example, if the original question asks about "the Manhattan project", in the follow up question, do not just say "the project", but use the full name "the Manhattan project". Your related questions must be in the same language as the original question.
+
+    Here are the contexts of the question:
+
+    {context}
+
+    Remember, based on the original question and related contexts, suggest three such further questions. Do NOT repeat the original question. Each related question should be no longer than 20 words. Here is the original question:
+    """.format(
+        context="\n\n".join([c["snippet"] for c in contexts])
+    )
+
+    try:
+        openai_client = new_async_client(_app)
+        llm_response = await openai_client.chat.completions.create(
+            model=_app.ctx.model,
+            messages=[
+                {"role": "system", "content": _lepton_more_questions_prompt},
+                {"role": "user", "content": query},
+            ],
+            max_tokens=512,
+            temperature=0.9,
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "ask_related_questions",
+                        "description": "Get a list of questions related to the original question and context.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "questions": {
+                                    "type": "array",
+                                    "items": {
+                                        "type": "string",
+                                        "description": "A related question to the original question and context.",
+                                    },
+                                }
+                            },
+                        },
+                        "required": ["questions"],
+                    },
+                }
+            ],
+            tool_choice="auto",
+        )
+        related = llm_response.choices[0].message.tool_calls[0].function.arguments
+        if isinstance(related, str):
+            related = json.loads(related)
+        logger.trace(f"Related questions: {related}")
+        return [{"question": _} for _ in related["questions"][:5]]
+    except Exception as e:
+        logger.error(
+            "Encountered error while generating related questions:"
+            f" {e}\n{traceback.format_exc()}"
+        )
+        return []
 
 async def _raw_stream_response(
     _app, contexts, llm_response, related_questions_future