Update app.py

app.py

@@ -1,76 +1,250 @@
 import os
 import warnings
-import gradio as gr
+import json
 from dotenv import load_dotenv
-
-from langchain.chains import ConversationalRetrievalChain
+from typing import Dict, Any, List, Optional
+import time
+from functools import lru_cache
+import logging
+ 
+ 
+ 
+from langchain.agents import Tool, AgentExecutor
+from langchain.tools.retriever import create_retriever_tool
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import AzureOpenAIEmbeddings
-from langchain_community.chat_models import AzureChatOpenAI
-
-# Patch Gradio bug
-import gradio_client.utils
-gradio_client.utils.json_schema_to_python_type = lambda schema, defs=None: "string"
-
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from openai import AzureOpenAI
+ 
 # Load environment variables
 load_dotenv()
 AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
 AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
-#AZURE_END_POINT_O3 = os.getenv("AZURE_END_POINT_O3")
 AZURE_OPENAI_LLM_DEPLOYMENT = os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT")
 AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")
-
-
-if not all([AZURE_OPENAI_API_KEY, 
-            AZURE_OPENAI_ENDPOINT,
-            #AZURE_END_POINT_O3, 
-            AZURE_OPENAI_LLM_DEPLOYMENT, 
-            AZURE_OPENAI_EMBEDDING_DEPLOYMENT]):
+ 
+if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_LLM_DEPLOYMENT, AZURE_OPENAI_EMBEDDING_DEPLOYMENT]):
     raise ValueError("Missing one or more Azure OpenAI environment variables.")
-
-# Suppress warnings
+ 
 warnings.filterwarnings("ignore")
-
-# Initialize Azure embedding model
+ 
+# Embeddings for retriever
 embeddings = AzureOpenAIEmbeddings(
     azure_deployment=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
     azure_endpoint=AZURE_OPENAI_ENDPOINT,
-    #azure_endpoint=AZURE_END_POINT_O3,
     openai_api_key=AZURE_OPENAI_API_KEY,
-    openai_api_version="2025-01-01-preview",  # updated to latest recommended version
+    openai_api_version="2025-01-01-preview",
     chunk_size=1000
 )
-
-# Load FAISS vector store
-vectorstore = FAISS.load_local(
-    "faiss_index_sysml", embeddings, allow_dangerous_deserialization=True
-)
-
-# Initialize Azure chat model
-llm = AzureChatOpenAI(
-    deployment_name=AZURE_OPENAI_LLM_DEPLOYMENT,
-    azure_endpoint=AZURE_OPENAI_ENDPOINT,
-    #azure_endpoint=AZURE_END_POINT_O3,
-    openai_api_key=AZURE_OPENAI_API_KEY,
-    openai_api_version="2025-01-01-preview",  # updated to latest recommended version
-    #temperature=0.5
-)
-
-# Build conversational RAG chain
-qa = ConversationalRetrievalChain.from_llm(
-    llm=llm,
-    retriever=vectorstore.as_retriever(),
-    return_source_documents=False
+ 
+# Get the directory where this script is located
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+ 
+# Build the absolute path to the faiss_index_sysml directory relative to this script
+FAISS_INDEX_PATH = os.path.join(SCRIPT_DIR, "faiss_index_sysml")
+# Load FAISS vectorstore
+vectorstore = FAISS.load_local(FAISS_INDEX_PATH, embeddings, allow_dangerous_deserialization=True)
+ 
+# Initialize Azure OpenAI client directly
+client = AzureOpenAI(
+    api_key=AZURE_OPENAI_API_KEY,
+    api_version="2025-01-01-preview",
+    azure_endpoint=AZURE_OPENAI_ENDPOINT
 )
+ 
+ 
+logger = logging.getLogger(__name__)
+# SysML retriever function
+@lru_cache(maxsize=100)
+def sysml_retriever(query: str) -> str:
+    start_time = time.time()
+    try:
+        results = vectorstore.similarity_search(query, k=100)
+        contexts = [doc.page_content for doc in results]
+        response = "\n\n".join(contexts)
+       
+        # Log performance metrics
+        duration = time.time() - start_time
+        print(f"Retrieval completed in {duration:.2f}s for query: {query[:50]}...")
+       
+        return response
+    except Exception as e:
+        logger.error(f"Retrieval error: {str(e)}")
+        return "Unable to retrieve information at this time."
+ 
+ 
+# sysml_retriever = create_retriever_tool(
+#     retriever=vectorstore.as_retriever(),
+#     name="SysMLRetriever",
+#     description="Use this to answer questions about SysML diagrams and modeling."
+# )
 
-history = []
-
-# Chatbot logic
+# Dummy functions
+def dummy_weather_lookup(location: str = "London") -> str:
+    return f"The weather in {location} is sunny and 25°C."
+ 
+def dummy_time_lookup(timezone: str = "UTC") -> str:
+    return f"The current time in {timezone} is 3:00 PM."
+ 
+# Tools definition for OpenAI function calling
+tools_definition = [
+    {
+        "type": "function",
+        "function": {
+            "name": "SysMLRetriever",
+            "description": "Use this to answer questions about SysML diagrams and modeling.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The search query to find information about SysML"
+                    }
+                },
+                "required": ["query"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "WeatherLookup",
+            "description": "Use this to look up the current weather in a specified location.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The location to look up the weather for"
+                    }
+                },
+                "required": ["location"]
+            }
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "TimeLookup",
+            "description": "Use this to look up the current time in a specified timezone.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "timezone": {
+                        "type": "string",
+                        "description": "The timezone to look up the current time for"
+                    }
+                },
+                "required": ["timezone"]
+            }
+        }
+    }
+]
+ 
+# Tool execution mapping
+tool_mapping = {
+    "SysMLRetriever": sysml_retriever,
+    "WeatherLookup": dummy_weather_lookup,
+    "TimeLookup": dummy_time_lookup
+}
+ 
+# Convert chat history
+def convert_history_to_messages(history):
+    messages = []
+    for user, bot in history:
+        messages.append({"role": "user", "content": user})
+        messages.append({"role": "assistant", "content": bot})
+    return messages
+ 
+# Main chatbot function with direct function calling
 def sysml_chatbot(message, history):
-    result = qa({"question": message, "chat_history": history})
-    answer = result["answer"]
-    history.append((message, answer))
-    return "", history
+    # Convert history to messages format
+    chat_messages = convert_history_to_messages(history)
+   
+    # Add system message at beginning
+    full_messages = [
+        {"role": "system", "content": "You are a helpful SysML modeling assistant and also a capable smart Assistant "}
+    ]
+    full_messages.extend(chat_messages)
+   
+    # Add current user message
+    full_messages.append({"role": "user", "content": message})
+   
+    try:
+        # First call to get either a direct answer or a function call
+        response = client.chat.completions.create(
+            model=AZURE_OPENAI_LLM_DEPLOYMENT,
+            messages=full_messages,
+            tools=tools_definition,
+            tool_choice={"type": "function", "function": {"name": "SysMLRetriever"}}
+        )
+       
+        assistant_message = response.choices[0].message
+       
+        # Check if the model wants to call a function
+        if assistant_message.tool_calls:
+            # Get the function call details
+            tool_call = assistant_message.tool_calls[0]
+            function_name = tool_call.function.name
+            function_args = json.loads(tool_call.function.arguments)
+            print("Attempting function calling...")
+            # Execute the function
+            if function_name in tool_mapping:
+                function_response = tool_mapping[function_name](**function_args)
+               
+                # Append the assistant's request and the function response to messages
+                full_messages.append({"role": "assistant", "content": None, "tool_calls": [
+                    {"id": tool_call.id, "type": "function", "function": {"name": function_name, "arguments": tool_call.function.arguments}}
+                ]})
+               
+                full_messages.append({
+                    "role": "tool",
+                    "tool_call_id": tool_call.id,
+                    "content": function_response
+                })
+               
+                # Second call to get the final answer based on the function result
+                second_response = client.chat.completions.create(
+                    model=AZURE_OPENAI_LLM_DEPLOYMENT,
+                    messages=full_messages
+                )
+               
+                answer = second_response.choices[0].message.content
+                print("Getting final response after function execution...")
+                print(f"Function '{function_name}' executed successfully. Response: {answer}")
+            else:
+                answer = f"I tried to use a function '{function_name}' that's not available. Let me try again with general knowledge: SysML is a modeling language for systems engineering that helps visualize and analyze complex systems."
+        else:
+            # Model provided a direct answer
+            answer = assistant_message.content
+       
+        history.append((message, answer))
+        return answer, history
+   
+    except Exception as e:
+        print(f"Error in function calling: {str(e)}")
+       
+        # Fallback to a direct response without function calling
+        try:
+            simple_messages = [
+                {"role": "system", "content": "You are a helpful SysML modeling assistant."}
+            ]
+            simple_messages.extend(chat_messages)
+            simple_messages.append({"role": "user", "content": message})
+           
+            fallback_response = client.chat.completions.create(
+                model=AZURE_OPENAI_LLM_DEPLOYMENT,
+                messages=simple_messages
+            )
+           
+            answer = fallback_response.choices[0].message.content
+        except Exception as fallback_error:
+            print(f"Error in fallback: {str(fallback_error)}")
+            answer = "I'm having trouble accessing my tools right now. SysML is a modeling language used in systems engineering to visualize and analyze complex systems through various diagram types."
+       
+        history.append((message, answer))
+        return answer, history
 
 # Gradio UI
 with gr.Blocks() as demo: