Update app.py

app.py

@@ -6,44 +6,57 @@ import os
 import re
 import time
 import uuid
-from typing import List, Dict, Any, Optional, Callable
+from typing import List, Dict, Any, Optional
 from dataclasses import dataclass
-from threading import Lock, Event, Thread
+from threading import Lock
+import threading
+import json
+import os
 import queue
 import traceback
+import uuid
+from typing import Coroutine, Dict, List, Any, Optional, Callable
+from dataclasses import dataclass
 from queue import Queue, Empty
+from threading import Lock, Event, Thread
+import threading
 from concurrent.futures import ThreadPoolExecutor
+import time
 import gradio as gr
 from openai import AsyncOpenAI, OpenAI
 import pyttsx3
 from rich.console import Console
 
-# --- Configuration ---
-BASE_URL = "http://localhost:1234/v1"
-BASE_API_KEY = "not-needed"
-# Using the sync client for the agent's internal sync calls
-CLIENT = OpenAI(base_url=BASE_URL, api_key=BASE_API_KEY)
-# Using the async client for the static async methods and instance methods
-BASE_CLIENT = AsyncOpenAI(base_url=BASE_URL, api_key=BASE_API_KEY)
-BASEMODEL_ID = "leroydyer/qwen/qwen3-0.6b-q4_k_m.gguf"
-
+BASE_URL="http://localhost:1234/v1"
+BASE_API_KEY="not-needed"
+BASE_CLIENT = AsyncOpenAI(
+    base_url=BASE_URL,
+    api_key=BASE_API_KEY
+) # Global state for client
+BASEMODEL_ID = "leroydyer/qwen/qwen3-0.6b-q4_k_m.gguf"  # Global state for selected model ID
+CLIENT =OpenAI(
+    base_url=BASE_URL,
+    api_key=BASE_API_KEY
+) # Global state for client
+# --- Global Variables (if needed) ---
 console = Console()
-
-# HuggingFace Spaces configuration (if needed)
+# --- Configuration ---
+LOCAL_BASE_URL = "http://localhost:1234/v1"
+LOCAL_API_KEY = "not-needed"
+# HuggingFace Spaces configuration
 HF_INFERENCE_URL = "https://api-inference.huggingface.co/models/"
 HF_API_KEY = os.getenv("HF_API_KEY", "")
-
-# Available model options (for UI reference, actual client is configured separately)
+# Available model options
 MODEL_OPTIONS = {
-    "Local LM Studio": BASE_URL,  # This is a URL, not a model ID
+    "Local LM Studio": LOCAL_BASE_URL,
     "Codellama 7B": "codellama/CodeLlama-7b-hf",
-    "Mistral 7B": "mistralai/Mistral-7B-v0.1",
+    "Mistral 7B": "mistralai/Mistral-7B-v0.1", 
     "Llama 2 7B": "meta-llama/Llama-2-7b-chat-hf",
     "Falcon 7B": "tiiuae/falcon-7b-instruct"
 }
-
 DEFAULT_TEMPERATURE = 0.7
 DEFAULT_MAX_TOKENS = 5000
+console = Console()
 
 # --- Canvas Artifact Support ---
 @dataclass
@@ -53,11 +66,7 @@ class CanvasArtifact:
     content: str
     title: str
     timestamp: float
-    metadata: Dict[str, Any] = None
-
-    def __post_init__(self):
-        if self.metadata is None:
-            self.metadata = {}
+    metadata: Dict[str, Any]
 
 @dataclass
 class LLMMessage:
@@ -67,7 +76,6 @@ class LLMMessage:
     conversation_id: str = None
     timestamp: float = None
     metadata: Dict[str, Any] = None
-
     def __post_init__(self):
         if self.message_id is None:
             self.message_id = str(uuid.uuid4())
@@ -81,7 +89,6 @@ class LLMRequest:
     message: LLMMessage
     response_event: str = None
     callback: Callable = None
-
     def __post_init__(self):
         if self.response_event is None:
             self.response_event = f"llm_response_{self.message.message_id}"
@@ -93,21 +100,18 @@ class LLMResponse:
     success: bool = True
     error: str = None
 
-# --- Event Manager ---
+# --- Event Manager (copied from your original code or imported) ---
 class EventManager:
     def __init__(self):
         self._handlers = defaultdict(list)
-        self._lock = Lock()
-
+        self._lock = threading.Lock()
     def register(self, event: str, handler: Callable):
         with self._lock:
             self._handlers[event].append(handler)
-
     def unregister(self, event: str, handler: Callable):
         with self._lock:
             if event in self._handlers and handler in self._handlers[event]:
                 self._handlers[event].remove(handler)
-
     def raise_event(self, event: str, data: Any):
         with self._lock:
             handlers = self._handlers[event][:]
@@ -118,7 +122,6 @@ class EventManager:
                 console.log(f"Error in event handler for {event}: {e}", style="bold red")
 
 EVENT_MANAGER = EventManager()
-
 def RegisterEvent(event: str, handler: Callable):
     EVENT_MANAGER.register(event, handler)
 
@@ -129,9 +132,9 @@ def UnregisterEvent(event: str, handler: Callable):
     EVENT_MANAGER.unregister(event, handler)
 
 class LLMAgent:
-    """Main Agent Driver !
-    Agent For Multiple messages at once ,
-    has a message queing service as well as a generator method for easy integration with console
+    """Main Agent Driver ! 
+    Agent For Multiple messages at once , 
+    has a message queing service as well as agenerator method for easy intergration with console     
     applications as well as ui !"""
     def __init__(
         self,
@@ -142,9 +145,9 @@ class LLMAgent:
         timeout: int = 30000,
         max_tokens: int = 5000,
         temperature: float = 0.3,
-        base_url: str = BASE_URL,
-        api_key: str = BASE_API_KEY,
-        generate_fn: Callable[[List[Dict[str, str]]], str] = None # Changed to sync function
+        base_url: str = "http://localhost:1234/v1",
+        api_key: str = "not-needed",
+        generate_fn: Callable[[List[Dict[str, str]]], Coroutine[Any, Any, str]] = None
     ):
         self.model_id = model_id
         self.system_prompt = system_prompt or "You are a helpful AI assistant."
@@ -154,25 +157,22 @@ class LLMAgent:
         self.is_running = False
         self._stop_event = Event()
         self.processing_thread = None
-
         # Conversation tracking
         self.conversations: Dict[str, List[LLMMessage]] = {}
         self.max_history_length = 20
-        # Use the provided generate function or the default sync one
-        self._generate = generate_fn or self._default_generate_sync
+        self._generate = generate_fn or self._default_generate
         self.api_key = api_key
-        self.base_url = base_url
+        self.base_url = base_url        
         self.max_tokens = max_tokens
         self.temperature = temperature
-        # Use the global async client for instance methods if needed
-        self.async_client = BASE_CLIENT
-
+        self.async_client = self.CreateClient(base_url, api_key)
         # Active requests waiting for responses
         self.pending_requests: Dict[str, LLMRequest] = {}
         self.pending_requests_lock = Lock()
 
-        # Canvas artifacts
-        self.canvas_artifacts: Dict[str, List[CanvasArtifact]] = defaultdict(list)
+        # Canvas Artifacts - NEW
+        self.canvas_artifacts: Dict[str, List[CanvasArtifact]] = {}
+        self.canvas_lock = Lock()
 
         # Register internal event handlers
         self._register_event_handlers()
@@ -186,7 +186,6 @@ class LLMAgent:
             console.log(f"[yellow]TTS not available: {e}[/yellow]")
             self.speech_enabled = False
         console.log("[bold green]🚀 Enhanced LLM Agent Initialized[/bold green]")
-
         # Start the processing thread immediately
         self.start()
 
@@ -210,19 +209,19 @@ class LLMAgent:
                 clean_text = re.sub(r'`.*?`', '', clean_text)
                 clean_text = clean_text.strip()
                 if clean_text:
-                    self.tts_engine.say(clean_text)
+                    self.tts_engine.say(clean_text) 
                     self.tts_engine.runAndWait()
                 else:
-                    self.tts_engine.say(text)
-                    self.tts_engine.runAndWait()
+                    self.tts_engine.say(text)  
+                    self.tts_engine.runAndWait()                    
             except Exception as e:
                 console.log(f"[red]TTS Error: {e}[/red]")
-        thread = Thread(target=_speak, daemon=True)
+        thread = threading.Thread(target=_speak, daemon=True)
         thread.start()
 
-    def _default_generate_sync(self, messages: List[Dict[str, str]]) -> str:
-        """Default sync generate function if none provided"""
-        return self._call_llm_sync(messages)
+    async def _default_generate(self, messages: List[Dict[str, str]]) -> str:
+        """Default generate function if none provided"""
+        return await self.openai_generate(messages)
 
     def _register_event_handlers(self):
         """Register internal event handlers for response routing"""
@@ -240,12 +239,10 @@ class LLMAgent:
             else:
                 console.log(f"No pending request found for: {response.request_id}", style="yellow")
                 return
-
         # Raise the specific response event
         if request.response_event:
             console.log(f"[bold green]Raising event: {request.response_event}[/bold green]")
             RaiseEvent(request.response_event, response)
-
         # Call callback if provided
         if request.callback:
             try:
@@ -287,10 +284,9 @@ class LLMAgent:
                 request.message
             )
             console.log(f"Calling LLM with {len(messages)} messages")
-            # Call LLM using the sync generate function
-            response_content = self._generate(messages)
-            console.log(f"[bold green]LLM response received: {response_content[:50]}...[/bold green]")
-
+            # Call LLM - Use sync call for thread compatibility
+            response_content = self._call_llm_sync(messages)
+            console.log(f"[bold green]LLM response received: {response_content}...[/bold green]")
             # Create response message
             response_message = LLMMessage(
                 role="assistant",
@@ -298,7 +294,6 @@ class LLMAgent:
                 conversation_id=request.message.conversation_id,
                 metadata={"request_id": request.message.message_id}
             )
-
             # Update conversation history
             self._add_to_conversation_history(
                 request.message.conversation_id or "default",
@@ -308,7 +303,6 @@ class LLMAgent:
                 request.message.conversation_id or "default",
                 response_message
             )
-
             # Create and send response
             response = LLMResponse(
                 message=response_message,
@@ -317,7 +311,6 @@ class LLMAgent:
             )
             console.log(f"[bold blue]Sending internal response for: {request.message.message_id}[/bold blue]")
             RaiseEvent("llm_internal_response", response)
-
         except Exception as e:
             console.log(f"[bold red]Error processing LLM request: {e}[/bold red]")
             traceback.print_exc()
@@ -352,7 +345,7 @@ class LLMAgent:
                 console.log(f"LLM call attempt {attempt + 1} failed: {e}")
                 if attempt == self.max_retries - 1:
                     raise e
-                time.sleep(1) # Wait before retry
+                 # Wait before retry
 
     def _process_queue(self):
         """Main queue processing loop"""
@@ -383,7 +376,6 @@ class LLMAgent:
         """Send a message to the LLM and get response via events"""
         if not self.is_running:
             raise RuntimeError("LLM Agent is not running. Call start() first.")
-
         # Create message
         message = LLMMessage(
             role=role,
@@ -391,19 +383,16 @@ class LLMAgent:
             conversation_id=conversation_id,
             metadata=metadata or {}
         )
-
         # Create request
         request = LLMRequest(
             message=message,
             response_event=response_event,
             callback=callback
         )
-
         # Store in pending requests BEFORE adding to queue
         with self.pending_requests_lock:
             self.pending_requests[message.message_id] = request
             console.log(f"Added to pending requests: {message.message_id}")
-
         # Add to queue
         try:
             self.request_queue.put(request, timeout=5.0)
@@ -424,14 +413,13 @@ class LLMAgent:
         # Create future for the response
         loop = asyncio.get_event_loop()
         response_future = loop.create_future()
-
         def chat_callback(response: LLMResponse):
             """Callback when LLM responds - thread-safe"""
             console.log(f"[bold yellow]✓ CHAT CALLBACK TRIGGERED![/bold yellow]")
             if not response_future.done():
                 if response.success:
                     content = response.message.content
-                    console.log(f"Callback received content: {content[:50]}...")
+                    console.log(f"Callback received content: {content}...")
                     # Schedule setting the future result on the main event loop
                     loop.call_soon_threadsafe(response_future.set_result, content)
                 else:
@@ -440,7 +428,6 @@ class LLMAgent:
                     loop.call_soon_threadsafe(response_future.set_result, error_msg)
             else:
                 console.log(f"[bold red]Future already done, ignoring callback[/bold red]")
-
         console.log(f"Sending message to LLM agent...")
         # Extract the actual message content from the messages list
         user_message = ""
@@ -450,7 +437,6 @@ class LLMAgent:
                 break
         if not user_message.strip():
             return ""
-
         # Send message with callback using the queue system
         try:
             message_id = self.send_message(
@@ -503,51 +489,11 @@ class LLMAgent:
         if conversation_id in self.conversations:
             del self.conversations[conversation_id]
 
-    # --- Canvas Methods ---
-    def add_artifact(self, conversation_id: str, artifact_type: str, content: str, title: str = "", metadata: Dict = None):
-        """Add an artifact to the canvas for a conversation."""
-        artifact = CanvasArtifact(
-            id=str(uuid.uuid4()),
-            type=artifact_type,
-            content=content,
-            title=title,
-            timestamp=time.time(),
-            metadata=metadata or {}
-        )
-        self.canvas_artifacts[conversation_id].append(artifact)
-
-    def get_canvas_artifacts(self, conversation_id: str = "default") -> List[CanvasArtifact]:
-        """Get all artifacts for a conversation."""
-        return self.canvas_artifacts.get(conversation_id, [])
-
-    def get_canvas_summary(self, conversation_id: str = "default") -> List[Dict[str, Any]]:
-        """Get a summary of artifacts for display."""
-        artifacts = self.get_canvas_artifacts(conversation_id)
-        return [{"id": a.id, "type": a.type, "title": a.title, "timestamp": a.timestamp} for a in artifacts]
-
-    def clear_canvas(self, conversation_id: str = "default"):
-        """Clear canvas artifacts for a conversation."""
-        if conversation_id in self.canvas_artifacts:
-            self.canvas_artifacts[conversation_id] = []
-
-    async def chat_with_canvas(self, user_message: str, conversation_id: str = "default", include_canvas: bool = False) -> str:
-        """
-        Chat method that can optionally include canvas content in the prompt.
-        """
-        messages = [{"role": "user", "content": user_message}]
-        if include_canvas:
-            canvas_artifacts = self.get_canvas_artifacts(conversation_id)
-            if canvas_artifacts:
-                canvas_content = "\n\n--- CANVAS CONTENT ---\n"
-                for artifact in canvas_artifacts:
-                    canvas_content += f"\n[{artifact.type}] {artifact.title or 'Untitled'}:\n{artifact.content}\n"
-                canvas_content += "\n--- END CANVAS CONTENT ---\n"
-                # Add canvas content as a system message
-                messages.insert(0, {"role": "system", "content": canvas_content})
-        return await self.chat(messages)
+    async def _chat(self, messages: List[Dict[str, str]]) -> str:
+        return await self._generate(messages)
 
     @staticmethod
-    async def openai_generate(messages: List[Dict[str, str]], max_tokens: int = 8096, temperature: float = 0.4, model: str = BASEMODEL_ID, tools=None) -> str:
+    async def openai_generate(messages: List[Dict[str, str]], max_tokens: int = 8096, temperature: float = 0.4, model: str = BASEMODEL_ID,tools=None) -> str:
         """Static method for generating responses using OpenAI API"""
         try:
             resp = await BASE_CLIENT.chat.completions.create(
@@ -563,6 +509,45 @@ class LLMAgent:
             console.log(f"[bold red]Error in openai_generate: {e}[/bold red]")
             return f"[LLM_Agent Error - openai_generate: {str(e)}]"
 
+    async def _call_(self, messages: List[Dict[str, str]]) -> str:
+        """Internal call method using instance client"""
+        try:
+            resp = await self.async_client.chat.completions.create(
+                model=self.model_id,
+                messages=messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens
+            )
+            response_text = resp.choices[0].message.content or ""
+            return response_text
+        except Exception as e:
+            console.log(f"[bold red]Error in _call_: {e}[/bold red]")
+            return f"[LLM_Agent Error - _call_: {str(e)}]"  
+
+    @staticmethod              
+    def CreateClient(base_url: str, api_key: str) -> AsyncOpenAI:
+        '''Create async OpenAI Client required for multi tasking'''
+        return AsyncOpenAI(
+            base_url=base_url,
+            api_key=api_key
+        ) 
+
+    @staticmethod
+    async def fetch_available_models(base_url: str, api_key: str) -> List[str]:
+        """Fetches available models from the OpenAI API."""
+        try:
+            async_client = AsyncOpenAI(base_url=base_url, api_key=api_key)
+            models = await async_client.models.list()
+            model_choices = [model.id for model in models.data]
+            return model_choices
+        except Exception as e:
+            console.log(f"[bold red]LLM_Agent Error fetching models: {e}[/bold red]")
+            return ["LLM_Agent Error fetching models"]   
+
+    def get_models(self) -> List[str]:             
+        """Get available models using instance credentials"""
+        return asyncio.run(self.fetch_available_models(self.base_url, self.api_key))
+
     def get_queue_size(self) -> int:
         """Get current queue size"""
         return self.request_queue.qsize()
@@ -582,208 +567,74 @@ class LLMAgent:
             "model": self.model_id
         }
 
-class AI_Agent:
-    def __init__(self, model_id: str, system_prompt: str = "You are a helpful assistant. Respond concisely in 1-2 sentences.", history: List[Dict] = None):
-        self.model_id = model_id
-        self.system_prompt = system_prompt
-        self.history = history or []
-        self.conversation_id = f"conv_{uuid.uuid4().hex[:8]}"
-        # Create agent instance - using the static async method as the generate function
-        self.client = LLMAgent(
-            model_id=model_id,
-            system_prompt=self.system_prompt,
-            generate_fn=lambda msgs: asyncio.run(LLMAgent.openai_generate(msgs, model=model_id))
-        )
-        console.log(f"[bold green]✓ MyAgent initialized with model: {model_id}[/bold green]")
-
-    async def call_llm(self, messages: List[Dict], use_history: bool = True) -> str:
-        """
-        Send messages to LLM and get response
-        Args:
-            messages: List of message dicts with 'role' and 'content'
-            use_history: Whether to include conversation history
-        Returns:
-            str: LLM response
-        """
-        try:
-            console.log(f"[bold yellow]Sending {len(messages)} messages to LLM (use_history: {use_history})...[/bold yellow]")
-            # Enhance messages based on history setting
-            enhanced_messages = await self._enhance_messages(messages, use_history)
-            response = await self.client.chat(enhanced_messages)
-            console.log(f"[bold green]✓ Response received ({len(response)} chars)[/bold green]")
-            # Update conversation history ONLY if we're using history
-            if use_history:
-                self._update_history(messages, response)
-            return response
-        except Exception as e:
-            console.log(f"[bold red]✗ ERROR: {e}[/bold red]")
-            traceback.print_exc()
-            return f"Error: {str(e)}"
-
-    async def _enhance_messages(self, messages: List[Dict], use_history: bool) -> List[Dict]:
-        """Enhance messages with system prompt and optional history"""
-        enhanced = []
-        # Add system prompt if not already in messages
-        has_system = any(msg.get('role') == 'system' for msg in messages)
-        if not has_system and self.system_prompt:
-            enhanced.append({"role": "system", "content": self.system_prompt})
-        # Add conversation history only if requested
-        if use_history and self.history:
-            enhanced.extend(self.history[-10:])  # Last 10 messages for context
-        # Add current messages
-        enhanced.extend(messages)
-        return enhanced
-
-    def _update_history(self, messages: List[Dict], response: str):
-        """Update conversation history with new exchange"""
-        # Add user messages to history
-        for msg in messages:
-            if msg.get('role') in ['user', 'assistant']:
-                self.history.append(msg)
-        # Add assistant response to history
-        self.history.append({"role": "assistant", "content": response})
-        # Keep history manageable (last 20 exchanges)
-        if len(self.history) > 40:  # 20 user + 20 assistant messages
-            self.history = self.history[-40:]
-
-    async def simple_query(self, query: str) -> str:
-        """Simple one-shot query method - NO history/context"""
-        messages = [{"role": "user", "content": query}]
-        return await self.call_llm(messages, use_history=False)
-
-    async def multi_turn_chat(self, user_input: str) -> str:
-        """Multi-turn chat that maintains context across calls"""
-        messages = [{"role": "user", "content": user_input}]
-        response = await self.call_llm(messages, use_history=True)
-        return response
-
-    def get_conversation_summary(self) -> Dict:
-        """Get conversation summary"""
-        return {
-            "conversation_id": self.conversation_id,
-            "total_messages": len(self.history),
-            "user_messages": len([msg for msg in self.history if msg.get('role') == 'user']),
-            "assistant_messages": len([msg for msg in self.history if msg.get('role') == 'assistant']),
-            "recent_exchanges": self.history[-4:] if self.history else []
-        }
-
-    def clear_history(self):
-        """Clear conversation history"""
-        self.history.clear()
-        console.log("[bold yellow]Conversation history cleared[/bold yellow]")
-
-    def update_system_prompt(self, new_prompt: str):
-        """Update the system prompt"""
-        self.system_prompt = new_prompt
-        console.log(f"[bold blue]System prompt updated[/bold blue]")
-
-    def stop(self):
-        """Stop the client gracefully"""
-        if hasattr(self, 'client') and self.client:
-            self.client.stop()
-            console.log("[bold yellow]MyAgent client stopped[/bold yellow]")
-
-    async def contextual_query(self, query: str, context_messages: List[Dict] = None,
-                             context_text: str = None, context_files: List[str] = None) -> str:
-        """
-        Query with specific context but doesn't update main history
-        Args:
-            query: The user question
-            context_messages: List of message dicts for context
-            context_text: Plain text context (will be converted to system message)
-            context_files: List of file paths to read and include as context
-        """
-        messages = []
-        # Add system prompt
-        if self.system_prompt:
-            messages.append({"role": "system", "content": self.system_prompt})
-        # Handle different context types
-        if context_messages:
-            messages.extend(context_messages)
-        if context_text:
-            messages.append({"role": "system", "content": f"Additional context: {context_text}"})
-        if context_files:
-            file_context = await self._read_files_context(context_files)
-            if file_context:
-                messages.append({"role": "system", "content": f"File contents:\n{file_context}"})
-        # Add the actual query
-        messages.append({"role": "user", "content": query})
-        return await self.call_llm(messages, use_history=False)
-
-    async def _read_files_context(self, file_paths: List[str]) -> str:
-        """Read multiple files and return as context string"""
-        contexts = []
-        for file_path in file_paths:
-            try:
-                if os.path.exists(file_path):
-                    with open(file_path, 'r', encoding='utf-8') as f:
-                        content = f.read()
-                        contexts.append(f"--- {os.path.basename(file_path)} ---\n{content}")
-                else:
-                    console.log(f"[bold yellow]File not found: {file_path}[/bold yellow]")
-            except Exception as e:
-                console.log(f"[bold red]Error reading file {file_path}: {e}[/bold red]")
-        return "\n".join(contexts) if contexts else ""
+    # --- ADDED CANVAS FUNCTIONALITY ---
+    def add_canvas_artifact(self, conversation_id: str, artifact_type: str, content: str, title: str = ""):
+        """Add an artifact to the canvas for a specific conversation."""
+        conv_id = conversation_id or "default"
+        with self.canvas_lock:
+            if conv_id not in self.canvas_artifacts:
+                self.canvas_artifacts[conv_id] = []
+            artifact = CanvasArtifact(
+                id=str(uuid.uuid4()),
+                type=artifact_type,
+                content=content,
+                title=title,
+                timestamp=time.time(),
+                metadata={}
+            )
+            self.canvas_artifacts[conv_id].append(artifact)
+        console.log(f"[green]Added {artifact_type} artifact to canvas '{conv_id}'[/green]")
+
+    def get_canvas_summary(self, conversation_id: str) -> List[Dict]:
+        """Get a summary of artifacts on the canvas for JSON display."""
+        conv_id = conversation_id or "default"
+        with self.canvas_lock:
+            artifacts = self.canvas_artifacts.get(conv_id, [])
+            # Convert artifacts to dictionaries for JSON serialization
+            return [
+                {
+                    "id": art.id,
+                    "type": art.type,
+                    "title": art.title,
+                    "timestamp": art.timestamp,
+                    "content_preview": art.content[:100] + "..." if len(art.content) > 100 else art.content
+                }
+                for art in artifacts
+            ]
+
+    def clear_canvas(self, conversation_id: str):
+        """Clear all artifacts from the canvas for a specific conversation."""
+        conv_id = conversation_id or "default"
+        with self.canvas_lock:
+            if conv_id in self.canvas_artifacts:
+                self.canvas_artifacts[conv_id].clear()
+        console.log(f"[yellow]Cleared canvas artifacts for '{conv_id}'[/yellow]")
+
+    async def chat_with_canvas(self, message: str, conversation_id: str, include_canvas: bool = False):
+        """Chat method that can optionally include canvas context."""
+        messages = [{"role": "user", "content": message}]
+        
+        if include_canvas:
+            artifacts = self.get_canvas_summary(conversation_id)
+            if artifacts:
+                canvas_context = "Current Canvas Context:\\n" + "\\n".join([
+                    f"- [{art['type'].upper()}] {art['title'] or 'Untitled'}: {art['content_preview']}"
+                    for art in artifacts
+                ])
+                messages.insert(0, {"role": "system", "content": canvas_context})
 
-    async def query_with_code_context(self, query: str, code_snippets: List[str] = None,
-                                    code_files: List[str] = None) -> str:
-        """
-        Specialized contextual query for code-related questions
-        """
-        code_context = "CODE CONTEXT:\n"
-        if code_snippets:
-            for i, snippet in enumerate(code_snippets, 1):
-                code_context += f"\nSnippet {i}:\n```\n{snippet}\n```\n"
-        if code_files:
-            # Read code files and include them
-            for file_path in code_files:
-                if file_path.endswith(('.py', '.js', '.java', '.cpp', '.c', '.html', '.css')):
-                    code_context += f"\nFile: {file_path}\n```\n"
-                    try:
-                        with open(file_path, 'r') as f:
-                            code_context += f.read()
-                    except Exception as e:
-                        code_context += f"Error reading file: {e}"
-                    code_context += "\n```\n"
-        return await self.contextual_query(query, context_text=code_context)
+        return await self.chat(messages)
 
-    async def multi_context_query(self, query: str, contexts: Dict[str, Any]) -> str:
-        """
-        Advanced contextual query with multiple context types
-        Args:
-            query: The user question
-            contexts: Dict with various context types
-                - 'messages': List of message dicts
-                - 'text': Plain text context
-                - 'files': List of file paths
-                - 'urls': List of URLs
-                - 'code': List of code snippets or files
-                - 'metadata': Any additional metadata
-        """
-        all_context_messages = []
-        # Build context from different sources
-        if contexts.get('text'):
-            all_context_messages.append({"role": "system", "content": f"Context: {contexts['text']}"})
-        if contexts.get('messages'):
-            all_context_messages.extend(contexts['messages'])
-        if contexts.get('files'):
-            file_context = await self._read_files_context(contexts['files'])
-            if file_context:
-                all_context_messages.append({"role": "system", "content": f"File Contents:\n{file_context}"})
-        if contexts.get('code'):
-            code_context = "\n".join([f"Code snippet {i}:\n```\n{code}\n```"
-                                    for i, code in enumerate(contexts['code'], 1)])
-            all_context_messages.append({"role": "system", "content": f"Code Context:\n{code_context}"})
-        if contexts.get('metadata'):
-            all_context_messages.append({"role": "system", "content": f"Metadata: {contexts['metadata']}"})
-        return await self.contextual_query(query, context_messages=all_context_messages)
 
+console = Console()
 
 # --- LCARS Styled Gradio Interface ---
 class LcarsInterface:
     def __init__(self):
-        # Start with the configured local client
+        # Start with HuggingFace by default for Spaces
+        self.use_huggingface = True
         self.agent = LLMAgent(generate_fn=LLMAgent.openai_generate)
+        self.current_conversation = "default"
 
     def create_interface(self):
         """Create the full LCARS-styled interface"""
@@ -885,10 +736,13 @@ class LcarsInterface:
             100% { opacity: 1; }
         }
         """
-
         with gr.Blocks(css=lcars_css, theme=gr.themes.Default(), title="LCARS Terminal") as interface:
             with gr.Column(elem_classes="lcars-container"):
                 # Header
+                with gr.Sidebar():
+                    gr.LoginButton() 
+
+
                 with gr.Row(elem_classes="lcars-header"):
                     gr.Markdown("""
                     <div style="text-align: center; width: 100%;">
@@ -906,11 +760,40 @@ class LcarsInterface:
                     with gr.Column(scale=1):
                         # Configuration Panel
                         with gr.Column(elem_classes="lcars-panel"):
+                            # Connection Type Selector
+                            with gr.Row(elem_classes="lcars-panel"):
+                            
+                                connection_type = gr.Radio(label = "### 🌐 CONNECTION TYPE",
+                                    choices=["HuggingFace Inference", "Local LM Studio"],
+                                    value="HuggingFace Inference",
+                                    elem_classes="lcars-input"
+                                )                            
                             gr.Markdown("### 🔧 CONFIGURATION")
+                            # Connection-specific settings
+                            with gr.Row(visible=False) as local_settings:
+                                base_url = gr.Textbox(
+                                    value=LOCAL_BASE_URL,
+                                    label="LM Studio URL",
+                                    elem_classes="lcars-input"
+                                )
+                                api_key = gr.Textbox(
+                                    value=LOCAL_API_KEY,
+                                    label="API Key",
+                                    type="password",
+                                    elem_classes="lcars-input"
+                                )
+                            with gr.Row(visible=True) as hf_settings:
+                                hf_api_key = gr.Textbox(
+                                    value=HF_API_KEY,
+                                    label="HuggingFace API Key",
+                                    type="password",
+                                    elem_classes="lcars-input",
+                                    placeholder="Get from https://huggingface.co/settings/tokens"
+                                )
                             with gr.Row():
                                 model_dropdown = gr.Dropdown(
-                                    choices=list(MODEL_OPTIONS.keys())[1:], # Exclude the 'Local LM Studio' URL entry
-                                    value=list(MODEL_OPTIONS.keys())[1], # Default to Codellama 7B
+                                    choices=list(MODEL_OPTIONS.keys())[1:],
+                                    value=list(MODEL_OPTIONS.keys())[1],
                                     label="AI Model",
                                     elem_classes="lcars-input"
                                 )
@@ -924,16 +807,16 @@ class LcarsInterface:
                         # Canvas Artifacts
                         with gr.Column(elem_classes="lcars-panel"):
                             gr.Markdown("### 🎨 CANVAS ARTIFACTS")
-                            artifact_display = gr.JSON(label="Canvas Summary")
+                            artifact_display = gr.JSON(label="")
                             with gr.Row():
                                 refresh_artifacts_btn = gr.Button("🔄 Refresh", elem_classes="lcars-button")
                                 clear_canvas_btn = gr.Button("🗑️ Clear Canvas", elem_classes="lcars-button")
                     # Main Content Area
                     with gr.Column(scale=2):
                         # Code Canvas
-                        with gr.Accordion("💻 COLLABORATIVE CODE CANVAS", open=True):
+                        with gr.Accordion("💻 COLLABORATIVE CODE CANVAS", open=False):
                             code_editor = gr.Code(
-                                value="# Welcome to LCARS Collaborative Canvas\nprint('Hello, Starfleet!')",
+                                value="# Welcome to LCARS Collaborative Canvas\\nprint('Hello, Starfleet!')",
                                 language="python",
                                 lines=15,
                                 label=""
@@ -945,107 +828,133 @@ class LcarsInterface:
                         # Chat Interface
                         with gr.Column(elem_classes="lcars-panel"):
                             gr.Markdown("### 💬 MISSION LOG")
-                            chatbot = gr.Chatbot(label="", height=300, elem_classes="lcars-chatbot")
+                            chatbot = gr.Chatbot(label="", height=300)
                             with gr.Row():
                                 message_input = gr.Textbox(
                                     placeholder="Enter your command or query...",
                                     show_label=False,
                                     lines=2,
-                                    elem_classes="lcars-input"
+                                    scale=4
                                 )
-                                send_btn = gr.Button("🚀 SEND", elem_classes="lcars-button")
+                                send_btn = gr.Button("🚀 SEND", elem_classes="lcars-button", scale=1)
                         # Status
                         with gr.Row():
                             status_display = gr.Textbox(
                                 value="LCARS terminal operational. Awaiting commands.",
                                 label="Status",
-                                max_lines=2,
-                                elem_classes="lcars-input"
+                                max_lines=2
                             )
                             with gr.Column(scale=0):
                                 clear_chat_btn = gr.Button("🗑️ Clear Chat", elem_classes="lcars-button")
                                 new_session_btn = gr.Button("🆕 New Session", elem_classes="lcars-button")
 
             # === EVENT HANDLERS ===
-            def update_agent_config(model_key, temp_val, max_tok_val, speech_enabled):
-                # Map UI model key to actual model ID
-                model_id = MODEL_OPTIONS.get(model_key, BASEMODEL_ID)
-                
-                # Update agent attributes
-                self.agent.model_id = model_id
-                self.agent.temperature = temp_val
-                self.agent.max_tokens = max_tok_val
-                self.agent.speech_enabled = speech_enabled
-                
-                # Update TTS if enabled/disabled
-                if speech_enabled and not self.agent.speech_enabled:
-                    try:
-                        self.agent.tts_engine = pyttsx3.init()
-                        self.agent.setup_tts()
-                        self.agent.speech_enabled = True
-                    except Exception as e:
-                        console.log(f"[yellow]TTS re-enable failed: {e}[/yellow]")
-                elif not speech_enabled and self.agent.speech_enabled:
-                    self.agent.speech_enabled = False
-                
-                return f"✅ Config updated: {model_key}, T={temp_val}, MaxTok={max_tok_val}, Speech={speech_enabled}"
-
-            def get_artifacts():
-                return self.agent.get_canvas_summary("default") # Assuming single conversation for UI
-
-            def clear_canvas():
-                self.agent.clear_canvas("default")
-                return [], "✅ Canvas cleared"
-
-            def clear_chat():
-                self.agent.clear_conversation("default")
-                return [], "✅ Chat cleared"
-
-            def new_session():
-                self.agent.clear_conversation("default")
-                self.agent.clear_canvas("default")
-                return [], "# New session started\nprint('Ready!')", "🆕 New session started", []
+            def switch_connection(connection_type):
+                if connection_type == "Local LM Studio":
+                    return [
+                        gr.update(visible=True),
+                        gr.update(visible=False),
+                        gr.update(choices=list(MODEL_OPTIONS.keys())[1:], value=list(MODEL_OPTIONS.keys())[1])
+                    ]
+                else:
+                    return [
+                        gr.update(visible=False),
+                        gr.update(visible=True),
+                        gr.update(choices=list(MODEL_OPTIONS.keys())[1:], value=list(MODEL_OPTIONS.keys())[1])
+                    ]
+
+            async def fetch_models_updated(connection_type, base_url_val, api_key_val, hf_api_key_val):
+                # Fixed: Removed the 'use_huggingface' parameter
+                if connection_type == "Local LM Studio":
+                    models = await LLMAgent.fetch_available_models(
+                        base_url_val, api_key_val
+                    )
+                else:
+                    # Using the HF_INFERENCE_URL and the key
+                    models = await LLMAgent.fetch_available_models(
+                        HF_INFERENCE_URL, hf_api_key_val
+                    )
+                if models:
+                    return gr.update(choices=models, value=models[0])
+                return gr.update(choices=["No models found"])
+
+            def update_agent_connection(connection_type, model_id, base_url_val, api_key_val, hf_api_key_val):
+                # Fixed: Removed the 'use_huggingface' parameter from the constructor
+                use_hf = connection_type == "HuggingFace Inference"
+                if use_hf:
+                    # Use the model_id directly (it's the model name like 'codellama/CodeLlama-7b-hf')
+                    self.agent = LLMAgent(
+                        model_id=model_id,
+                        base_url=HF_INFERENCE_URL,
+                        api_key=hf_api_key_val,
+                        generate_fn=LLMAgent.openai_generate
+                    )
+                    return f"✅ Switched to HuggingFace: {model_id}"
+                else:
+                    self.agent = LLMAgent(
+                        model_id=model_id, 
+                        base_url=base_url_val,
+                        api_key=api_key_val,
+                        generate_fn=LLMAgent.openai_generate
+                    )
+                    return f"✅ Switched to Local: {base_url_val}"
 
             async def process_message(message, history, speech_enabled):
                 if not message.strip():
-                    return "", history, "Please enter a message", self.agent.get_canvas_summary("default")
+                    return "", history, "Please enter a message"
                 history = history + [[message, None]]
                 try:
-                    # For simplicity, use the basic chat method here. Canvas integration can be added if needed.
-                    response = await self.agent.chat([{"role": "user", "content": message}])
+                    # Fixed: Uses the new chat_with_canvas method which includes canvas context
+                    response = await self.agent.chat_with_canvas(
+                        message, self.current_conversation, include_canvas=True
+                    )
                     history[-1][1] = response
                     if speech_enabled and self.agent.speech_enabled:
                         self.agent.speak(response)
-                    artifacts = self.agent.get_canvas_summary("default")
+                    artifacts = self.agent.get_canvas_summary(self.current_conversation)
                     status = f"✅ Response received. Canvas artifacts: {len(artifacts)}"
                     return "", history, status, artifacts
                 except Exception as e:
                     error_msg = f"❌ Error: {str(e)}"
                     history[-1][1] = error_msg
-                    return "", history, error_msg, self.agent.get_canvas_summary("default")
+                    return "", history, error_msg, self.agent.get_canvas_summary(self.current_conversation)
+
+            def get_artifacts():
+                return self.agent.get_canvas_summary(self.current_conversation)
+
+            def clear_canvas():
+                self.agent.clear_canvas(self.current_conversation)
+                return [], "✅ Canvas cleared"
+
+            def clear_chat():
+                self.agent.clear_conversation(self.current_conversation)
+                return [], "✅ Chat cleared"
+
+            def new_session():
+                self.agent.clear_conversation(self.current_conversation)
+                self.agent.clear_canvas(self.current_conversation)
+                return [], "# New session started\\nprint('Ready!')", "🆕 New session started", []
 
             # Connect events
-            update_config_btn.click(
-                update_agent_config,
-                inputs=[model_dropdown, temperature, max_tokens, speech_toggle],
-                outputs=status_display
-            )
+            connection_type.change(switch_connection, inputs=connection_type, 
+                                 outputs=[local_settings, hf_settings, model_dropdown])
+            fetch_models_btn.click(fetch_models_updated, 
+                                 inputs=[connection_type, base_url, api_key, hf_api_key],
+                                 outputs=model_dropdown)
+            update_config_btn.click(update_agent_connection,
+                                inputs=[connection_type, model_dropdown, base_url, api_key, hf_api_key],
+                                outputs=status_display)
+            send_btn.click(process_message,
+                         inputs=[message_input, chatbot, speech_toggle],
+                         outputs=[message_input, chatbot, status_display, artifact_display])
+            message_input.submit(process_message,
+                               inputs=[message_input, chatbot, speech_toggle],
+                               outputs=[message_input, chatbot, status_display, artifact_display])
             refresh_artifacts_btn.click(get_artifacts, outputs=artifact_display)
             clear_canvas_btn.click(clear_canvas, outputs=[artifact_display, status_display])
             clear_chat_btn.click(clear_chat, outputs=[chatbot, status_display])
             new_session_btn.click(new_session, outputs=[chatbot, code_editor, status_display, artifact_display])
-            send_btn.click(
-                process_message,
-                inputs=[message_input, chatbot, speech_toggle],
-                outputs=[message_input, chatbot, status_display, artifact_display]
-            )
-            message_input.submit(
-                process_message,
-                inputs=[message_input, chatbot, speech_toggle],
-                outputs=[message_input, chatbot, status_display, artifact_display]
-            )
             interface.load(get_artifacts, outputs=artifact_display)
-
         return interface
 
 # --- Main Application ---
@@ -1056,10 +965,11 @@ def main():
         console.log("[green]🌐 Detected HuggingFace Space[/green]")
     else:
         console.log("[blue]💻 Running locally[/blue]")
-
     interface = LcarsInterface()
     demo = interface.create_interface()
-    demo.launch(share=is_space)
+    demo.launch(
+        share=is_space
+    )
 
 if __name__ == "__main__":
     main()