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 AsyncQdrantClient
from qdrant_client.models import PointStruct, Distance, VectorParams
from chainlit.data.sql_alchemy import SQLAlchemyDataLayer

# --- CONFIGURAZIONE DATABASE (PostgreSQL) ---
# Assicurati che user/password coincidano con il tuo docker-compose.yml
# Sintassi: postgresql+asyncpg://user:password@host:port/dbname
DATABASE_URL = "postgresql+asyncpg://user:password@postgres:5432/ai_station"

# Attiviamo il salvataggio su DB
cl.data_layer = SQLAlchemyDataLayer(conn_string=DATABASE_URL)
# -

# --- CONFIGURAZIONE HARD-CODED ---
OLLAMA_URL = "http://192.168.1.243:11434"
# ---------------------------------

USER_ROLES = {
    'moglie@esempio.com': 'business',
    'ingegnere@esempio.com': 'engineering',
    'architetto@esempio.com': 'architecture',
    'admin@esempio.com': 'admin'
}

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 = AsyncQdrantClient(url="http://qdrant:6333")
    collection_name = "documents"
    
    # Check if collection exists
    if not await client.collection_exists(collection_name):
        await client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(size=768, distance=Distance.COSINE)
        )
    
    return client

async def get_embeddings(text):
    client = ollama.Client(host=OLLAMA_URL)
    
    # Limite di sicurezza per evitare errori 500 su Ollama
    limit = 2000 
    if len(text) > limit:
        text = text[:limit]
    
    try:
        response = client.embed(model='nomic-embed-text', input=text)
        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)
        
        if not query_embedding:
            return ""
            
        # Usa query_points (nuova API per AsyncClient)
        search_result = await qdrant_client.query_points(
            collection_name="documents",
            query=query_embedding,
            limit=3
        )
        
        hits = search_result.points
        contexts = []
        if hits:
            for hit in hits:
                try:
                    if hit.payload:
                        # --- FIX IMPORTANTE: Recupera il contenuto reale ---
                        nome_file = hit.payload.get('file_name', 'Sconosciuto')
                        contenuto = hit.payload.get('content', '') 
                        
                        contexts.append(f"--- Documento: {nome_file} ---\n{contenuto}\n----------------")
                except Exception as e:
                    print(f"Error parsing hit: {e}")
        
        return "\n".join(contexts)
    except Exception as e:
        print(f"Errore ricerca Qdrant: {e}")
        return ""

@cl.on_chat_start
async def chat_start():
    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)
    
    welcome_msg = f"Welcome, {user_role.capitalize()}!"
    await cl.Message(content=welcome_msg).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.Client(host=OLLAMA_URL)
        history = cl.user_session.get("history", [])
        history = limit_history(history)
        
        # --- PASSO 1: Gestione Upload e Indexing (PRIMA della ricerca) ---
        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'):
                        # Encoding utf-8 per sicurezza
                        with open(dest_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                        
                        embeddings = await get_embeddings(content)
                        if embeddings:
                            qdrant_client = await connect_to_qdrant()
                            point_id = str(uuid.uuid4())
                            
                            # --- FIX IMPORTANTE: Salviamo anche il contenuto nel payload ---
                            point = PointStruct(
                                id=point_id, 
                                vector=embeddings, 
                                payload={
                                    "file_name": element.name,
                                    "content": content 
                                }
                            )
                            await 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()

        # --- PASSO 2: Cerca nei documenti ---
        context_text = await search_qdrant(message.content, user_role)
        
        if context_text:
            system_prompt = f"Usa esclusivamente il seguente contesto per rispondere alla domanda. Se la risposta non è nel contesto, dillo.\n\nContesto:\n{context_text}"
            history.insert(0, {"role": "system", "content": system_prompt})
        
        history.append({"role": "user", "content": message.content})
        
        # --- PASSO 3: Chat Generation ---
        response = client.chat(model='qwen2.5-coder:7b', messages=history)
        
        # Code Extraction
        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()