Update app.py

app.py

@@ -1,7 +1,10 @@
+"""
+Main FastAPI application integrating all components with Hugging Face Inference Endpoint.
+"""
 import gradio as gr
 import fastapi
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import HTMLResponse, FileResponse
+from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
 from fastapi import FastAPI, Request, Form, UploadFile, File
 import os
 import time
@@ -10,6 +13,13 @@ import json
 import shutil
 import uvicorn
 from pathlib import Path
+from typing import Dict, List, Optional, Any
+import io
+import numpy as np
+from scipy.io.wavfile import write
+
+# Import our modules
+from local_llm import run_llm, run_llm_with_memory, clear_memory, get_memory_sessions, get_model_info, test_endpoint
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -31,47 +41,19 @@ if html_template.exists() and not static_html.exists():
 # Mount static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Mock data and functions to simulate the real implementation
-SESSIONS = {}
-
-def generate_session_id():
-    """Generate a unique session ID."""
-    import uuid
-    return str(uuid.uuid4())
-
+# Helper functions for mock implementations
 def mock_transcribe(audio_bytes):
     """Mock function to simulate speech-to-text."""
-    # In production, this would use Whisper
     logger.info("Transcribing audio...")
-    time.sleep(1)  # Simulate processing time
+    time.sleep(0.5)  # Simulate processing time
     return "This is a mock transcription of the audio."
 
-def mock_agent_response(text, session_id="default"):
-    """Mock function to simulate agent reasoning."""
-    # In production, this would use a real LLM
-    logger.info(f"Processing query: {text}")
-    time.sleep(1.5)  # Simulate processing time
-    
-    # Simple keyword-based responses
-    if "5g" in text.lower():
-        return "5G is the fifth generation of cellular networks, offering higher speeds, lower latency, and more capacity than previous generations."
-    elif "telecom" in text.lower():
-        return "Telecommunications (telecom) refers to the exchange of information over significant distances by electronic means."
-    elif "webrtc" in text.lower():
-        return "WebRTC (Web Real-Time Communication) is a free, open-source project that enables web browsers and mobile applications to have real-time communication via simple APIs."
-    else:
-        return "I'm an AI assistant specialized in telecom topics. Feel free to ask me about 5G, network technologies, or telecommunications in general."
-
 def mock_synthesize_speech(text):
     """Mock function to simulate text-to-speech."""
-    # In production, this would use a real TTS engine
     logger.info("Synthesizing speech...")
     time.sleep(0.5)  # Simulate processing time
     
     # Create a dummy audio file
-    import numpy as np
-    from scipy.io.wavfile import write
-    
     sample_rate = 22050
     duration = 2  # seconds
     t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
@@ -83,9 +65,6 @@ def mock_synthesize_speech(text):
     with open(output_file, "rb") as f:
         audio_bytes = f.read()
     
-    # Clean up
-    os.remove(output_file)
-    
     return audio_bytes
 
 # Routes for the API
@@ -94,6 +73,15 @@ async def root():
     """Serve the main UI."""
     return FileResponse("static/index.html")
 
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    endpoint_status = test_endpoint()
+    return {
+        "status": "ok",
+        "endpoint": endpoint_status
+    }
+
 @app.post("/api/transcribe")
 async def transcribe(file: UploadFile = File(...)):
     """Transcribe audio to text."""
@@ -103,17 +91,24 @@ async def transcribe(file: UploadFile = File(...)):
         return {"transcription": text}
     except Exception as e:
         logger.error(f"Transcription error: {str(e)}")
-        return {"error": f"Failed to transcribe audio: {str(e)}"}
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Failed to transcribe audio: {str(e)}"}
+        )
 
 @app.post("/api/query")
 async def query_agent(input_text: str = Form(...), session_id: str = Form("default")):
     """Process a text query with the agent."""
     try:
-        response = mock_agent_response(input_text, session_id)
+        response = run_llm_with_memory(input_text, session_id=session_id)
+        logger.info(f"Query: {input_text[:30]}... Response: {response[:30]}...")
         return {"response": response}
     except Exception as e:
         logger.error(f"Query error: {str(e)}")
-        return {"error": f"Failed to process query: {str(e)}"}
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Failed to process query: {str(e)}"}
+        )
 
 @app.post("/api/speak")
 async def speak(text: str = Form(...)):
@@ -127,19 +122,135 @@ async def speak(text: str = Form(...)):
         )
     except Exception as e:
         logger.error(f"Speech synthesis error: {str(e)}")
-        return {"error": f"Failed to synthesize speech: {str(e)}"}
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Failed to synthesize speech: {str(e)}"}
+        )
 
 @app.post("/api/session")
 async def create_session():
     """Create a new session."""
-    session_id = generate_session_id()
-    SESSIONS[session_id] = {"created_at": time.time()}
+    import uuid
+    session_id = str(uuid.uuid4())
+    clear_memory(session_id)
     return {"session_id": session_id}
 
+@app.delete("/api/session/{session_id}")
+async def delete_session(session_id: str):
+    """Delete a session."""
+    success = clear_memory(session_id)
+    if success:
+        return {"message": f"Session {session_id} cleared"}
+    else:
+        return JSONResponse(
+            status_code=404,
+            content={"error": f"Session {session_id} not found"}
+        )
+
+@app.get("/api/sessions")
+async def list_sessions():
+    """List all active sessions."""
+    return {"sessions": get_memory_sessions()}
+
+@app.get("/api/model_info")
+async def model_info():
+    """Get information about the model."""
+    return get_model_info()
+
+@app.post("/api/complete")
+async def complete_flow(
+    request: Request,
+    audio_file: UploadFile = File(None),
+    text_input: str = Form(None),
+    session_id: str = Form("default")
+):
+    """
+    Complete flow: audio to text to agent to speech.
+    """
+    try:
+        # If audio file provided, transcribe it
+        if audio_file:
+            audio_bytes = await audio_file.read()
+            text_input = mock_transcribe(audio_bytes)
+            logger.info(f"Transcribed input: {text_input[:30]}...")
+        
+        # Process with agent
+        if not text_input:
+            return JSONResponse(
+                status_code=400,
+                content={"error": "No input provided"}
+            )
+        
+        response = run_llm_with_memory(text_input, session_id=session_id)
+        logger.info(f"Agent response: {response[:30]}...")
+        
+        # Synthesize speech
+        audio_bytes = mock_synthesize_speech(response)
+        
+        # Save audio to a temporary file
+        import tempfile
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        temp_file.write(audio_bytes)
+        temp_file.close()
+        
+        # Generate URL for audio
+        host = request.headers.get("host", "localhost")
+        scheme = request.headers.get("x-forwarded-proto", "http")
+        audio_url = f"{scheme}://{host}/audio/{os.path.basename(temp_file.name)}"
+        
+        return {
+            "input": text_input,
+            "response": response,
+            "audio_url": audio_url
+        }
+    
+    except Exception as e:
+        logger.error(f"Complete flow error: {str(e)}")
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Failed to process: {str(e)}"}
+        )
+
+@app.get("/audio/{filename}")
+async def get_audio(filename: str):
+    """
+    Serve temporary audio files.
+    """
+    try:
+        # Ensure filename only contains safe characters
+        import re
+        if not re.match(r'^[a-zA-Z0-9_.-]+$', filename):
+            return JSONResponse(
+                status_code=400,
+                content={"error": "Invalid filename"}
+            )
+        
+        temp_dir = tempfile.gettempdir()
+        file_path = os.path.join(temp_dir, filename)
+        
+        if not os.path.exists(file_path):
+            return JSONResponse(
+                status_code=404,
+                content={"error": "File not found"}
+            )
+        
+        return FileResponse(
+            file_path,
+            media_type="audio/wav",
+            filename=filename
+        )
+    
+    except Exception as e:
+        logger.error(f"Audio serving error: {str(e)}")
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Failed to serve audio: {str(e)}"}
+        )
+
 # Gradio interface
 with gr.Blocks(title="AGI Telecom POC", css="footer {visibility: hidden}") as interface:
     gr.Markdown("# AGI Telecom POC Demo")
-    gr.Markdown("This is a demonstration of the AGI Telecom Proof of Concept. The full interface is available via the direct API.")
+    gr.Markdown("This is a demonstration of the AGI Telecom Proof of Concept using a Hugging Face Inference Endpoint.")
     
     with gr.Row():
         with gr.Column():
@@ -165,10 +276,13 @@ with gr.Blocks(title="AGI Telecom POC", css="footer {visibility: hidden}") as in
             
             # Status and info
             status_output = gr.Textbox(label="Status", value="Ready")
+            endpoint_status = gr.Textbox(label="Endpoint Status", value="Checking endpoint connection...")
     
     # Link components with functions
     def update_session():
-        new_id = generate_session_id()
+        import uuid
+        new_id = str(uuid.uuid4())
+        clear_memory(new_id)
         status = f"Created new session: {new_id}"
         return new_id, status
     
@@ -189,7 +303,7 @@ with gr.Blocks(title="AGI Telecom POC", css="footer {visibility: hidden}") as in
             text = mock_transcribe(audio_bytes)
             
             # Get response
-            response = mock_agent_response(text, session)
+            response = run_llm_with_memory(text, session)
             
             # Synthesize
             audio_bytes = mock_synthesize_speech(response)
@@ -210,7 +324,7 @@ with gr.Blocks(title="AGI Telecom POC", css="footer {visibility: hidden}") as in
     )
     
     query_btn.click(
-        lambda text, session: mock_agent_response(text, session),
+        lambda text, session: run_llm_with_memory(text, session),
         inputs=[text_input, session_id],
         outputs=[response_output]
     )
@@ -227,6 +341,16 @@ with gr.Blocks(title="AGI Telecom POC", css="footer {visibility: hidden}") as in
         inputs=[audio_input, session_id],
         outputs=[transcription_output, response_output, audio_output, status_output]
     )
+    
+    # Check endpoint on load
+    def check_endpoint():
+        status = test_endpoint()
+        if status["status"] == "connected":
+            return f"✅ Connected to endpoint: {status['message']}"
+        else:
+            return f"❌ Error connecting to endpoint: {status['message']}"
+    
+    gr.on_load(lambda: gr.update(value=check_endpoint()), outputs=endpoint_status)
 
 # Mount Gradio app
 app = gr.mount_gradio_app(app, interface, path="/gradio")