ai-station/app.py

import os
import chainlit as cl
import re
from datetime import datetime
import shutil
import uuid
import ollama
from qdrant_client import QdrantClient
from qdrant_client.http.models import PointStruct

# --- CONFIGURAZIONE HARD-CODED PER ROMPERE IL BLOCCO 127.0.0.1 ---
OLLAMA_URL = "http://192.168.1.243:11434"
# -----------------------------------------------------------------------------

# Define user roles mapping
USER_ROLES = {
    'moglie@esempio.com': 'business',
    'ingegnere@esempio.com': 'engineering',
    'architetto@esempio.com': 'architecture',
    'admin@esempio.com': 'admin'
}

# Define the path for workspaces
WORKSPACES_DIR = "./workspaces"

def create_workspace(user_role):
    workspace_path = os.path.join(WORKSPACES_DIR, user_role)
    if not os.path.exists(workspace_path):
        os.makedirs(workspace_path)

def save_code_to_file(code, user_role):
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    file_name = f"code_{timestamp}.py"
    file_path = os.path.join(WORKSPACES_DIR, user_role, file_name)
    
    with open(file_path, "w") as file:
        file.write(code)
    
    return file_path

def limit_history(history):
    if len(history) > 20:
        history = history[-20:]
    return history

async def connect_to_qdrant():
    client = QdrantClient("http://qdrant:6333")
    collection_name = "documents"
    
    try:
        client.get_collection(collection_name)
    except Exception:
        client.create_collection(
            collection_name=collection_name,
            vectors_config={"size": 768, "distance": "Cosine"}
        )
    
    return client

async def get_embeddings(text):
    # --- FIX: Splitto Host e Port per evitare confusione ---
    client = ollama.Client(host=OLLAMA_URL) # Uso l'URL intero
    
    # Controllo lunghezza testo
    if len(text) > 12000:
        text = text[:12000]
    
    try:
        response = client.embed(model='nomic-embed-text', input=text)
        
        # Gestione risposta
        if 'embeddings' in response:
            return response['embeddings'][0]
        return response.get('embedding')
    except Exception as e:
        print(f"Errore Embedding: {e}")
        return []

async def search_qdrant(query_text, user_role):
    """Cerca documenti pertinenti su Qdrant"""
    try:
        qdrant_client = await connect_to_qdrant()
        query_embedding = await get_embeddings(query_text)
        
        # Se non trova embedding (errore connessione), non cercare
        if not query_embedding:
            return ""
            
        # Cerca i 3 documenti più simili
        search_result = qdrant_client.search(
            collection_name="documents",
            query_vector=query_embedding,
            limit=3
        )
        
        contexts = []
        # FIX: controllo sicurezza per evitare 'list index out of range'
        if search_result:
            for hit in search_result:
                try:
                    if 'payload' in hit and 'file_name' in hit['payload']:
                        contexts.append(f"Documento: {hit['payload']['file_name']}")
                except Exception:
                    pass
        
        return "\n".join(contexts)
    except Exception as e:
        print(f"Errore ricerca Qdrant: {e}")
        return ""

@cl.on_chat_start
async def chat_start():
    # Hardcode per test
    user_email = "admin@esempio.com"
    user_role = USER_ROLES.get(user_email, 'guest')
    
    create_workspace(user_role)
    
    cl.user_session.set("history", [])
    cl.user_session.set("role", user_role)
    
    if user_role == 'admin':
        await cl.Message(content="Welcome, Admin!").send()
    elif user_role == 'engineering':
        await cl.Message(content="Welcome, Engineer!").send()
    elif user_role == 'business':
        await cl.Message(content="Welcome, Business User!").send()
    elif user_role == 'architecture':
        await cl.Message(content="Welcome, Architect!").send()
    else:
        await cl.Message(content="Welcome, Guest!").send()

@cl.on_message
async def message(message):
    user_role = cl.user_session.get("role", 'guest')
    
    if not user_role:
        await cl.Message(content="User role not found").send()
        return
    
    try:
        # Client Ollama URL Hardcoded
        client = ollama.Client(host=OLLAMA_URL)
        
        # History & Sliding Window
        history = cl.user_session.get("history", [])
        history = limit_history(history)
        
        # --- RAG STEP 1: Cerca nei documenti ---
        context_text = await search_qdrant(message.content, user_role)
        
        # Se trova contesto, iniettalo
        if context_text:
            system_prompt = f"Contexto dai documenti:\n{context_text}\n\nRispondi usando questo contesto."
            history.insert(0, {"role": "system", "content": system_prompt})
        
        history.append({"role": "user", "content": message.content})
        
        # Gestione Upload e Indexing
        if message.elements:
            uploaded_files = []
            for element in message.elements:
                try:
                    dest_path = os.path.join(WORKSPACES_DIR, user_role, element.name)
                    with open(element.path, 'rb') as src, open(dest_path, 'wb') as dst:
                        shutil.copyfileobj(src, dst)
                    
                    if element.name.endswith('.txt'):
                        with open(dest_path, 'r') as f:
                            content = f.read()
                        
                        # Indexing
                        embeddings = await get_embeddings(content)
                        if embeddings:
                            qdrant_client = await connect_to_qdrant()
                            point_id = uuid.uuid4()
                            point = PointStruct(id=point_id, vector=embeddings, payload={"file_name": element.name})
                            qdrant_client.upsert(collection_name="documents", points=[point])
                            await cl.Message(content=f"Documento '{element.name}' indicizzato.").send()
                    
                    uploaded_files.append(element.name)
                except Exception as e:
                    await cl.Message(content=f"Error saving {element.name}: {e}").send()
            
            if uploaded_files:
                 await cl.Message(content=f"Files saved: {', '.join(uploaded_files)}").send()
        
        # Chat
        response = client.chat(model='qwen2.5-coder:7b', messages=history)
        
        # Code Extracting
        code_blocks = re.findall(r"```python(.*?)```", response['message']['content'], re.DOTALL)
        
        elements = []
        if code_blocks:
            for code in code_blocks:
                file_path = save_code_to_file(code, user_role)
                elements.append(cl.File(name=os.path.basename(file_path), path=file_path))
        
        history.append({"role": "assistant", "content": response['message']['content']})
        cl.user_session.set("history", history)
        
        await cl.Message(content=response['message']['content'], elements=elements).send()
        
    except Exception as e:
        await cl.Message(content=f"Error: {e}").send()