Update app.py

app.py

@@ -3,38 +3,32 @@ import gradio as gr
 import warnings
 import json
 from dotenv import load_dotenv
-from typing import Dict, Any, List, Optional
+from typing import List
 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_core.messages import HumanMessage, AIMessage, SystemMessage
 from openai import AzureOpenAI
 
-
 # Patch Gradio bug
 import gradio_client.utils
 gradio_client.utils.json_schema_to_python_type = lambda schema, defs=None: "string"
 
-
 # Load environment variables
 load_dotenv()
 AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
 AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
 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_OPENAI_LLM_DEPLOYMENT, AZURE_OPENAI_EMBEDDING_DEPLOYMENT]):
     raise ValueError("Missing one or more Azure OpenAI environment variables.")
- 
+
 warnings.filterwarnings("ignore")
- 
-# Embeddings for retriever
+
+# Embeddings
 embeddings = AzureOpenAIEmbeddings(
     azure_deployment=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
     azure_endpoint=AZURE_OPENAI_ENDPOINT,
@@ -42,57 +36,41 @@ embeddings = AzureOpenAIEmbeddings(
     openai_api_version="2025-01-01-preview",
     chunk_size=1000
 )
- 
-# Get the directory where this script is located
+
+# Vectorstore
 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
+
+# OpenAI client
 client = AzureOpenAI(
     api_key=AZURE_OPENAI_API_KEY,
     api_version="2025-01-01-preview",
     azure_endpoint=AZURE_OPENAI_ENDPOINT
 )
- 
- 
+
+# Logger
 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
+        return "\n\n".join(contexts)
     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."
-# )
 
 # 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 for function calling
 tools_definition = [
     {
         "type": "function",
@@ -102,10 +80,7 @@ tools_definition = [
             "parameters": {
                 "type": "object",
                 "properties": {
-                    "query": {
-                        "type": "string",
-                        "description": "The search query to find information about SysML"
-                    }
+                    "query": {"type": "string", "description": "The search query to find information about SysML"}
                 },
                 "required": ["query"]
             }
@@ -119,14 +94,11 @@ tools_definition = [
             "parameters": {
                 "type": "object",
                 "properties": {
-                    "location": {
-                        "type": "string",
-                        "description": "The location to look up the weather for"
-                    }
+                    "location": {"type": "string", "description": "The location to look up the weather for"}
                 },
                 "required": ["location"]
             }
-        },
+        }
     },
     {
         "type": "function",
@@ -136,24 +108,21 @@ tools_definition = [
             "parameters": {
                 "type": "object",
                 "properties": {
-                    "timezone": {
-                        "type": "string",
-                        "description": "The timezone to look up the current time for"
-                    }
+                    "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 = []
@@ -162,113 +131,72 @@ def convert_history_to_messages(history):
         messages.append({"role": "assistant", "content": bot})
     return messages
 
-
-history = []
-
-
-# Main chatbot function with direct function calling
+# Chatbot logic
 def sysml_chatbot(message, 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})
-   
+        {"role": "system", "content": "You are a helpful SysML modeling assistant and also a capable smart Assistant"}
+    ] + chat_messages + [{"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": "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."
+                answer = f"I tried to use a function '{function_name}' that's not available."
         else:
-            # Model provided a direct answer
             answer = assistant_message.content
-       
         history.append((message, answer))
-        return answer, history
-   
+        return "", 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
+        history.append((message, "Sorry, something went wrong."))
+        return "", history
 
-# Gradio UI
+# === Gradio UI ===
 with gr.Blocks(css="""
-    #submit-btn {
-        height: 100%;
-        background-color: #48CAE4;
-        color: white;
-        font-size: 1.5em;
-    }
+#submit-btn {
+    height: 100%;
+    background-color: #48CAE4;
+    color: white;
+    font-size: 1.5em;
+}
 """) as demo:
-    
+
     gr.Markdown("## SysModeler Chatbot")
 
     chatbot = gr.Chatbot(height=600)
-    
     with gr.Row():
         with gr.Column(scale=5):
             msg = gr.Textbox(
@@ -277,19 +205,14 @@ with gr.Blocks(css="""
                 show_label=False
             )
         with gr.Column(scale=1, min_width=50):
-            submit_btn = gr.Button("➤")
+            submit_btn = gr.Button("➤", elem_id="submit-btn")
 
     clear = gr.Button("Clear")
+    state = gr.State([])
 
-    state = gr.State(history)
-
-    submit_btn.click(fn=sysml_chatbot, inputs=[msg, state], outputs=["", chatbot])
-    msg.submit(fn=sysml_chatbot, inputs=[msg, state], outputs=["", chatbot])  # still supports enter key
-    
-    # submit_btn.click(fn=sysml_chatbot, inputs=[msg, state], outputs=[gr.Textbox.update(value=""), chatbot])
-    # msg.submit(fn=sysml_chatbot, inputs=[msg, state], outputs=[gr.Textbox.update(value=""), chatbot])
-
+    submit_btn.click(fn=sysml_chatbot, inputs=[msg, state], outputs=[msg, chatbot])
+    msg.submit(fn=sysml_chatbot, inputs=[msg, state], outputs=[msg, chatbot])
     clear.click(fn=lambda: ([], ""), inputs=None, outputs=[chatbot, msg])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()