Add tool-calling knowledge retrieval in AI responses

Enable function-calling flow so the assistant can query bank knowledge before producing final answers, and align config defaults for model and prompt handling.

Made-with: Cursor

ai_service.py

@@ -1,70 +1,117 @@
 import re
-from config import pc, index, EMBED_MODEL, HF_MODEL, PROMPT, hf_client
+import json
+import os
+from config import pc, index, EMBED_MODEL, hf_client, PROMPT
 from database import db_manager
 
+
+MODEL_NAME = "dolphin-mistral-24b-venice-edition"
+
 def clean_ai_response(text: str):
+    if not text: return ""
     cleaned_text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
     return cleaned_text.strip()
 
-async def get_ai_response(user_query: str, telegram_id: int = None):
-
-    if not pc or not index or not hf_client:
-        return "Ai service is not available at the moment. Please try again later."
-
-    # Save user message if database is available and telegram_id is provided
-    conversation_history = ""
-    if telegram_id and db_manager:
-        db_manager.save_message(telegram_id, user_query, "user")
-        conversation_history = db_manager.get_formatted_history(telegram_id, limit=6)        
-
+async def search_bank_knowledge(query: str):
     query_embedding = pc.inference.embed(
         model=EMBED_MODEL,
-        inputs=[user_query],
+        inputs=[query],
         parameters={"input_type": "query"}
     )
-
-    # Search Pinecone for Bank Context
+    
     search_results = index.query(
         vector=query_embedding[0].values,
         top_k=3,
         include_metadata=True
     )
     
-    retrieved_context = "\n".join([res.metadata['original_text'] for res in search_results.matches])
+    return "\n".join([res.metadata['original_text'] for res in search_results.matches])
 
-    
-    # Prepare the user content with conversation history and context
-    user_content = f"""
-        ### Historical Conversation:
-        {conversation_history}
+# تعريف الأداة (Tool) الخاصة بالبحث في وثائق البنك
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "search_bank_knowledge",
+            "description": "استخدم هذه الأداة للبحث في الملف التعريفي الرسمي لبنك حضرموت للحصول على معلومات دقيقة حول الخدمات، الهيكل التنظيمي، رأس المال، والسياسات.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "جملة البحث باللغة العربية (مثال: 'ما هو رأس مال بنك حضرموت؟' أو 'خدمات الأفراد')."
+                    }
+                },
+                "required": ["query"]
+            }
+        }
+    }
+]
 
-        ### Retrieved Context from Bank Documents:
-        {retrieved_context}
+async def get_ai_response(user_query: str, telegram_id: int = None):
+    if not pc or not index or not hf_client:
+        return "عذراً، خدمة الذكاء الاصطناعي غير متوفرة حالياً."
 
-        ### Current User Message:
-        {user_query}
+    # 1. إدارة تاريخ المحادثة
+    conversation_history = []
+    if telegram_id and db_manager:
+        db_manager.save_message(telegram_id, user_query, "user")
+        # جلب آخر 6 رسائل لتوفير السياق للموديل
+        raw_history = db_manager.get_history(telegram_id, limit=6)
+        for msg in raw_history:
+            conversation_history.append({"role": msg['role'], "content": msg['content']})
+    else:
+        conversation_history.append({"role": "user", "content": user_query})
 
-        Based on the above information, provide an accurate and helpful response to the customer:
-    """
-    print("User content:", user_content)
+    
+    messages = [{"role": "system", "content": PROMPT}] + conversation_history
 
-    completion = hf_client.chat.completions.create(
-        model=HF_MODEL,
-        messages=[
-            {"role": "system", "content": PROMPT},
-            {"role": "user", "content": user_content},
-        ],
-        temperature=0.1,
-        max_tokens=800,
-        top_p=0.9,
+    
+    response = hf_client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=messages,
+        tools=TOOLS,
+        tool_choice="auto",
+        temperature=0.1 
     )
 
-    ai_response = completion.choices[0].message.content
-    cleaned_response = clean_ai_response(ai_response)
+    response_message = response.choices[0].message
+    tool_calls = getattr(response_message, 'tool_calls', None)
 
-    # Save assistant response if database is available and telegram_id is provided
+    
+    if tool_calls:
+        
+        for tool_call in tool_calls:
+            function_args = json.loads(tool_call.function.arguments)
+            search_query = function_args.get("query")
+            
+            
+            extracted_context = await search_bank_knowledge(search_query)
+            
+            
+            messages.append(response_message)
+            messages.append({
+                "role": "tool",
+                "tool_call_id": tool_call.id,
+                "name": "search_bank_knowledge",
+                "content": extracted_context
+            })
+
+        
+        final_response = hf_client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            temperature=0.3
+        )
+        ai_final_content = final_response.choices[0].message.content
+    else:
+        
+        ai_final_content = response_message.content
+
+    
+    cleaned_response = clean_ai_response(ai_final_content)
+    
     if telegram_id and db_manager:
         db_manager.save_message(telegram_id, cleaned_response, "assistant")
     
-    return cleaned_response
-
+    return cleaned_response

config.py

@@ -20,15 +20,9 @@ TELEGRAM_URL = f"https://149.154.167.220/bot{TELEGRAM_TOKEN}/sendMessage" if TEL
 EMBED_MODEL = os.environ.get("EMBED_MODEL", "multilingual-e5-large")
 HF_MODEL = os.environ.get(
     "HF_MODEL",
-    "dphn/Dolphin-Mistral-24B-Venice-Edition:featherless-ai",
-)
-PROMPT = os.environ.get(
-    "PROMPT",
-    "You are a helpful customer service assistant for Hadhramout Bank. "
-    "Answer the user's question based on the provided context. If the context "
-    "doesn't contain the answer, politely say you don't have enough information "
-    "to help with that specific query."
+    "dphn/Dolphin-Mistral-24B-Venice-Edition",
 )
+PROMPT = os.environ.get("PROMPT")
 
 # Initialize clients only if API keys are available
 pc = None