SynAccel-AI/retrieval.py at main · Arkwinger/SynAccel-AI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
from pathlib import Path

from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
import chromadb


model = SentenceTransformer("all-MiniLM-L6-v2")

client = chromadb.PersistentClient(path="chroma_db")


def get_collection():

    try:
        return client.get_collection("documents")

    except Exception:
        return client.create_collection("documents")


def load_documents():

    docs_folder = Path("docs")

    try:
        client.delete_collection("documents")
    except Exception:
        pass

    collection = client.create_collection(
        name="documents"
    )

    doc_id = 0

    for file in docs_folder.iterdir():

        content = ""

        if file.suffix.lower() == ".pdf":

            try:

                reader = PdfReader(file)

                for page in reader.pages:

                    text = page.extract_text()

                    if text:
                        content += text + "\n"

            except Exception:
                continue

        elif file.suffix.lower() == ".txt":

            try:

                with open(file, "r", encoding="utf-8") as f:
                    content = f.read()

            except Exception:
                continue

        if not content.strip():
            continue

        # Better chunking with overlap
        chunk_size = 500
        overlap = 200

        chunks = []

        for i in range(0, len(content), chunk_size - overlap):
            chunks.append(content[i:i + chunk_size])

        for chunk in chunks:

            embedding = model.encode(chunk).tolist()

            collection.add(
                ids=[str(doc_id)],
                documents=[chunk],
                embeddings=[embedding],
                metadatas=[
                    {
                        "source": file.name
                    }
                ]
            )

            doc_id += 1


def search_documents(query):

    collection = get_collection()

    query_embedding = model.encode(query).tolist()

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=3,
        include=[
            "documents",
            "metadatas",
            "distances"
        ]
    )

    if not results:
        return None

    if not results.get("distances"):
        return None

    try:

        best_distance = results["distances"][0][0]

        print(f"\nBest match distance: {best_distance}")

        # More forgiving threshold
        if best_distance > 2.0:

            print("Match rejected - distance too high")

            return None

    except Exception:

        return None

    print("\n====================")
    print("QUERY:", query)

    for docs in results["documents"]:

        for doc in docs:

            print("\n--- MATCH ---")
            print(doc[:500])

    print("====================\n")

    return results