demeter

 1import uuid
 2from qdrant_client import models
 3from fastembed import TextEmbedding
 4from Qdrant.Client import client 
 5from openai import AzureOpenAI
 6import os
 7from dotenv import load_dotenv
 8
 9load_dotenv()
10
11class ResearcherAgent:
12    def __init__(self):
13        self.client = client
14        self.llm = AzureOpenAI(
15            api_key=os.getenv("AZURE_OPENAI_API_KEY"),
16            api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-12-01-preview"),
17            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
18        )
19        self.deployment_name = os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "gpt-4.1")
20        self.collection = "Knowledge_Base"
21        # FastEmbed is lightweight and runs locally on CPU
22        self.encoder = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
23        self._ensure_collection()
24
25    def _ensure_collection(self):
26        """Creates the text collection if it doesn't exist."""
27        if not self.client.collection_exists(self.collection):
28            self.client.create_collection(
29                collection_name=self.collection,
30                vectors_config=models.VectorParams(
31                    size=384, 
32                    distance=models.Distance.COSINE
33                )
34            )
35            print(f"📚 Created knowledge base: {self.collection}")
36
37    def ingest_text(self, text: str, source: str = "Manual"):
38        """Saves a chunk of text (e.g., from a PDF) into memory."""
39        # Embed the text
40        embedding = list(self.encoder.embed([text]))[0]
41        
42        # Upload to Qdrant
43        self.client.upsert(
44            collection_name=self.collection,
45            points=[
46                models.PointStruct(
47                    id=str(uuid.uuid4()),
48                    vector=embedding.tolist(),
49                    payload={"text": text, "source": source}
50                )
51            ]
52        )
53
54    def search(self, query: str, limit: int = 1) -> str:
55        """Retrieves relevant textbook pages and formats them as a string."""
56        # 1. Convert query to vector
57        query_vec = list(self.encoder.embed([query]))[0]
58        
59        # 2. Search Qdrant
60        response = self.client.query_points(
61            collection_name=self.collection,
62            query=query_vec,
63            limit=limit,
64            with_payload=True 
65        )
66
67        hits = response.points
68
69        # 3. Format as a single string (better for LLM Context)
70        if not hits:
71            return "No specific scientific manuals found for this issue."
72            
73        return "\n\n".join([
74            f"[Source: {hit.payload.get('source', 'Unknown')}]\n{hit.payload.get('text', '')}" 
75            for hit in hits
76        ])