yashasvi0409/Document_Intelligence

1import faiss2import numpy as np3 4# ✅ Import your embedding function (same file you already use in main.py)5from embeddings import embed_texts6 7 8class VectorStore:9    def __init__(self, dimension: int):10        # Inner Product index (works well if you normalize embeddings)11        self.index = faiss.IndexFlatIP(dimension)12        self.text_chunks = []13        self.metadata = []14 15    def add(self, embeddings, chunks, sources):16        """17        embeddings: list/np.array of shape (N, dim)18        chunks: list[str] length N19        sources: list[Any] length N20        """21        emb = np.array(embeddings, dtype="float32")22 23        # FAISS expects 2D array24        if emb.ndim == 1:25            emb = emb.reshape(1, -1)26 27        self.index.add(emb)28        self.text_chunks.extend(chunks)29        self.metadata.extend(sources)30 31    def search(self, query_embedding, top_k: int = 3):32        """33        query_embedding: np.array shape (1, dim) or (dim,)34        returns list of dicts with text/source/score35        """36        q = np.array(query_embedding, dtype="float32")37 38        # FAISS expects 2D array: (1, dim)39        if q.ndim == 1:40            q = q.reshape(1, -1)41 42        scores, indices = self.index.search(q, top_k)43 44        results = []45        for idx, score in zip(indices[0], scores[0]):46            if idx == -1:47                continue48            if idx >= len(self.text_chunks):49                continue50 51            results.append({52                "text": self.text_chunks[idx],53                "source": self.metadata[idx],54                "score": float(score)55            })56 57        return results58 59    # ✅ NEW: This is what routes.py should call60    def query(self, question: str, top_k: int = 3):61        """62        Converts question -> embedding, searches FAISS,63        then returns output in your UI format.64        """65        question = (question or "").strip()66        if not question:67            return {68                "answer": "Please ask a question.",69                "confidence": "LOW",70                "similarity_score": 0.0,71                "source_documents": []72            }73 74        # ✅ Get embedding for the question (embed_texts expects list[str])75        q_embs = embed_texts([question])76 77        # embed_texts may return numpy array or list78        q_embs = np.array(q_embs, dtype="float32")79 80        # Ensure shape (1, dim)81        if q_embs.ndim == 1:82            q_embs = q_embs.reshape(1, -1)83 84        # ✅ Search85        hits = self.search(q_embs[0], top_k=top_k)86 87        if not hits:88            return {89                "answer": "No relevant information found in the uploaded documents.",90                "confidence": "LOW",91                "similarity_score": 0.0,92                "source_documents": []93            }94 95        # ✅ Build answer from best chunk (top hit)96        best = hits[0]97        answer_text = best.get("text", "—")98        best_score = float(best.get("score", 0.0))99 100        # ✅ Simple confidence mapping101        if best_score >= 0.75:102            conf = "HIGH"103        elif best_score >= 0.45:104            conf = "MEDIUM"105        else:106            conf = "LOW"107 108        sources = []109        for h in hits:110            src = h.get("source")111            if src and src not in sources:112                sources.append(src)113 114        return {115            "answer": answer_text,116            "confidence": conf,117            "similarity_score": best_score,118            "source_documents": sources119        }