added regex

app.py

@@ -7,15 +7,8 @@ from pinecone import Pinecone, ServerlessSpec
 from groq import Groq
 from tqdm.auto import tqdm
 import streamlit as st
+import re
 
-# Required imports
-import json
-import time
-import os
-from sentence_transformers import SentenceTransformer
-from pinecone import Pinecone, ServerlessSpec
-from groq import Groq
-from tqdm.auto import tqdm
 
 # Constants (hardcoded)
 FILE_PATH = "anjibot_chunks.json"
@@ -61,10 +54,45 @@ for i in tqdm(range(0, len(data['id']), BATCH_SIZE)):
     to_upsert = list(zip(batch["id"], embeds, batch["metadata"]))
     index.upsert(vectors=to_upsert)
 
-def get_docs(query: str, top_k: int) -> list[str]:
-    xq = encoder.encode(query)
-    res = index.query(vector=xq.tolist(), top_k=top_k, include_metadata=True)
-    return [x["metadata"]['content'] for x in res["matches"]]
+def extract_course_code(text) -> list[str]:
+    pattern = r'\b(?:geds?|stats?|maths?|cosc|seng|itgy)\s*\d{3}\b'
+    match = re.findall(pattern, text, re.IGNORECASE)
+    return match if match else None
+    
+def get_docs(query: str, top_k: int, batch_size: int = 5, threshold: float = 0.66) -> list[str]:
+    queried_course_codes = extract_course_code(query)
+    
+    i = 0 
+    relevant_docs = []
+
+    while True:
+        xq = encoder.encode(query)
+        res = index.query(vector=xq.tolist(), top_k=batch_size, include_metadata=True, offset=i)
+
+        if len(res["matches"]) == 0:
+            break
+            
+        for match in res["matches"]:
+            similarity_score = match['score']
+            content = match["metadata"]['content']
+
+            if similarity_score >= threshold:
+                if queried_course_codes:
+                    for course_code in queried_course_codes:
+                        if course_code in content:
+                            relevant_docs.append(content)
+                            break  
+
+        if relevant_docs:
+            break
+
+        i += batch_size
+
+    if relevant_docs:
+        return relevant_docs
+    else:
+        return ["No exact match found for the course code, even after searching with a higher similarity score."]
+
 
 def get_response(query: str, docs: list[str]) -> str:
     system_message = (
@@ -92,7 +120,7 @@ def get_response(query: str, docs: list[str]) -> str:
 def handle_query(user_query: str):
 
     # Get relevant documents
-    docs = get_docs(user_query, top_k=5)
+    docs = get_docs(user_query, top_k=5, threshold=0.66)
 
     # Generate and return response
     response = get_response(user_query, docs=docs)