Graph Tool Integration

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
-from langchain.prompts import ChatPromptTemplate
-from langchain.schema import HumanMessage, SystemMessage, AIMessage
-from huggingface_hub import InferenceClient
+from graph_tool import generate_plot
 from metrics import EduBotMetrics
+from together import Together
+from together.types.chat.completion import CompletionChunk
 import os
 import time
 import logging
+import json
 import re
 
 # --- Environment and Logging Setup ---
@@ -18,13 +19,34 @@ if not hf_token:
     logger.warning("Neither HF_TOKEN nor HUGGINGFACEHUB_API_TOKEN is set, the application may not work.")
 
 # --- LLM Configuration ---
-client = InferenceClient(
-    provider="together",
-    api_key=hf_token,
-)
+client = Together(api_key=hf_token)
 
 metrics_tracker = EduBotMetrics(save_file="edu_metrics.json")
 
+# --- Tools ---
+
+tools = [
+    {
+        "name": "create_graph",
+        "description": "Generates a plot (bar, line, or pie) and returns it as an HTML-formatted Base64-encoded image string. The data and labels arguments must be JSON-encoded strings. Use simple, descriptive labels for plots. Use this tool to produce plots for practice questions, data visualization use cases, or any other case where the student may benefit from a graph.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "data_json": {"type": "string"},
+                "labels_json": {"type": "string"},
+                "plot_type": {"type": "string", "enum": ["bar", "line", "pie"]},
+                "title": {"type": "string"},
+                "x_label": {"type": "string"},
+                "y_label": {"type": "string"},
+            },
+            "required": ["data_json", "labels_json", "plot_type", "title"],
+        },
+    }
+]
+
+# --- BIND THE TOOL TO THE MODEL ---
+model_with_tools = llm.bind_tools([create_graph])
+
 # --- LLM Templates ---
 # Enhanced base system message
 SYSTEM_MESSAGE = """You are EduBot, an expert multi-concept tutor designed to facilitate genuine learning and understanding. Your primary mission is to guide students through the learning process rather than providing direct answers to academic work.
@@ -148,76 +170,99 @@ def smart_truncate(text, max_length=3000):
         return ' '.join(words[:-1]) + "... [Response truncated - ask for continuation]"
 
 def respond_with_enhanced_streaming(message, history):
-    """Streams the bot's response, handling errors with metrics tracking."""
-
-    # Start metrics timing
+    """Streams the bot's response, handling tool calls and errors with metrics tracking."""
     timing_context = metrics_tracker.start_timing()
     error_occurred = False
     error_message = None
-    response = ""
+    full_response = ""
     
     try:
-        # Build conversation history (last 5 exchanges)
         api_messages = [{"role": "system", "content": SYSTEM_MESSAGE}]
         if history:
+            # Prepare history for API call
             for exchange in history[-5:]:
-                if exchange.get("role") == "user":
-                    api_messages.append({"role": "user", "content": exchange["content"]})
-                elif exchange.get("role") == "assistant":
-                    api_messages.append({"role": "assistant", "content": exchange["content"]})
+                api_messages.append({"role": "user", "content": exchange[0]})
+                api_messages.append({"role": "assistant", "content": exchange[1]})
         
-        # Add current user message
         api_messages.append({"role": "user", "content": message})
-
-        # Mark provider API start
+        
         metrics_tracker.mark_provider_start(timing_context)
         
-        completion = client.chat.completions.create(
+        stream = client.chat.completions.create(
             model="Qwen/Qwen2.5-7B-Instruct",
             messages=api_messages,
             max_tokens=4096,
             temperature=0.7,
             top_p=0.9,
+            stream=True,
+            tools=tools, # Pass the tool definitions here
         )
+        
+        # Buffers to handle multi-chunk tool calls
+        tool_call_name = ""
+        tool_call_args_str = ""
+
+        for chunk in stream:
+            if isinstance(chunk, CompletionChunk):
+                # Handle text chunks
+                if chunk.choices and chunk.choices[0].delta.content:
+                    text_chunk = chunk.choices[0].delta.content
+                    full_response += text_chunk
+                    yield full_response
+                
+                # Handle tool call chunks
+                if chunk.choices and chunk.choices[0].delta.tool_calls:
+                    tool_call_delta = chunk.choices[0].delta.tool_calls[0]
+                    
+                    # Accumulate name and arguments from stream
+                    if tool_call_delta.function.name:
+                        tool_call_name = tool_call_delta.function.name
+                    if tool_call_delta.function.arguments:
+                        tool_call_args_str += tool_call_delta.function.arguments
+                    
+                    # Check if we have received the full tool call
+                    # This is a simple heuristic and might need refinement
+                    if tool_call_name and '}' in tool_call_args_str:
+                        try:
+                            tool_args = json.loads(tool_call_args_str)
+                            if tool_call_name == "create_graph":
+                                logger.info(f"Executing tool: {tool_call_name} with args: {tool_args}")
+                                graph_html = generate_plot(**tool_args)
+                                full_response += graph_html
+                                yield full_response
+                            
+                            # Reset buffers
+                            tool_call_name = ""
+                            tool_call_args_str = ""
+
+                        except json.JSONDecodeError:
+                            logger.error("JSON parsing failed for tool arguments.")
+                            # Yield an error or handle it silently
+                            full_response += f"<p style='color:red;'>Error parsing graph data.</p>"
+                            yield full_response
+                        except Exception as e:
+                            logger.exception("Error executing tool")
+                            full_response += f"<p style='color:red;'>Error executing tool: {e}</p>"
+                            yield full_response
 
-        # Mark provider API end
         metrics_tracker.mark_provider_end(timing_context)
+        logger.info(f"Response completed. Length: {len(full_response)} characters")
         
-        response = smart_truncate(completion.choices[0].message.content, max_length=3000)
-        
-        # Stream fake-chunks word by word
-        words = response.split()
-        partial_response = ""
-        
-        for i, word in enumerate(words):
-            partial_response += word + " "
-            if i % 4 == 0:
-                metrics_tracker.record_chunk(timing_context)
-                yield partial_response
-                time.sleep(0.03)
-        
-        logger.info(f"Response completed. Length: {len(response)} characters")
-        
-        metrics_tracker.record_chunk(timing_context)
-        yield response
-
     except Exception as e:
         error_occurred = True
         error_message = str(e)
         logger.exception("Error in response generation")
-        yield "Sorry, I encountered an error while generating the response."
+        yield "Sorry, an error occurred while generating the response."
 
     finally:
         metrics_tracker.log_interaction(
             query=message,
-            response=response,
+            response=full_response,
             timing_context=timing_context,
             error_occurred=error_occurred,
             error_message=error_message,
         )
 
-        
-
 # ===============================================================================
 # UI CONFIGURATION SECTION - ALL UI RELATED CODE CENTRALIZED HERE
 # ===============================================================================