server/tests/test_rag.py

import asyncio
import numpy as np
from app.services.vector_db import VectorDB
from app.services.vector_operations import VectorOperations
from app.services.ai_engine import AIEngine

async def test_rag():
    print("Testing RAG Pipeline...")
    
    # 1. Test Vector Search
    vector_db = VectorDB()
    website_id = 8
    
    if vector_db.load(website_id):
        print(f"✓ VectorDB loaded for website {website_id}")
        
        query = "tell me about open transport"
        print(f"\nQuery: {query}")
        
        embedding = await VectorOperations.get_embedding(query)
        query_vector = np.array(embedding, dtype=np.float32)
        
        results = vector_db.search(query_vector, k=5)
        print(f"✓ Found {len(results)} results")
        
        for i, (meta, score) in enumerate(results):
            print(f"  {i+1}. Score: {score:.4f} - {meta['content'][:100]}...")
            
        # 2. Test AI Engine with RAG Context
        print("\nTesting AI Engine Generation...")
        ai_engine = AIEngine()
        
        context = [r[0]['content'] for r in results]
        response, confidence, needs_contact = await ai_engine.generate_response(
            query=query,
            context=context,
            website_id=website_id
        )
        
        print(f"\nResponse: {response}")
        print(f"Confidence: {confidence}")
        print(f"Needs Contact: {needs_contact}")
        
    else:
        print(f"✗ Failed to load VectorDB for website {website_id}")

if __name__ == "__main__":
    asyncio.run(test_rag())