feat: Implement retriever and guardrail modules

guards/input.py

@@ -0,0 +1,30 @@
+import re
+
+def is_valid(query: str) -> bool:
+    """
+    Validates the user's query.
+    """
+    # Check for query length
+    if len(query) > 500:
+        return False
+
+    # Check for SQL injection patterns
+    sql_injection_patterns = [
+        r"(\s*(--|#|;))",
+        r"(\s*(union|select|insert|update|delete|drop|alter)\s+)",
+    ]
+    for pattern in sql_injection_patterns:
+        if re.search(pattern, query, re.IGNORECASE):
+            return False
+
+    return True
+
+if __name__ == '__main__':
+    # Example usage
+    valid_query = "What are the computer science courses?"
+    invalid_query_long = "a" * 501
+    invalid_query_sql = "SELECT * FROM students;"
+
+    print(f"'{valid_query}' is valid: {is_valid(valid_query)}")
+    print(f"'long query' is valid: {is_valid(invalid_query_long)}")
+    print(f"'{invalid_query_sql}' is valid: {is_valid(invalid_query_sql)}")

guards/output.py

@@ -0,0 +1,27 @@
+import re
+
+def sanitize(response: str) -> str:
+    """
+    Sanitizes the LLM's response.
+    """
+    # Redact email addresses
+    response = re.sub(r"[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+", "[REDACTED EMAIL]", response)
+
+    # Check for hallucination patterns
+    hallucination_patterns = [
+        r"as an ai language model",
+        r"i cannot answer that question",
+    ]
+    for pattern in hallucination_patterns:
+        if re.search(pattern, response, re.IGNORECASE):
+            return "I'm sorry, I don't have enough information to answer that question."
+
+    return response
+
+if __name__ == '__main__':
+    # Example usage
+    valid_response = "Professor Smith's email is john.smith@university.edu."
+    hallucination_response = "As an AI language model, I cannot provide that information."
+
+    print(f"Original: '{valid_response}'\nSanitized: '{sanitize(valid_response)}'")
+    print(f"\nOriginal: '{hallucination_response}'\nSanitized: '{sanitize(hallucination_response)}'")

rag/retriever.py

@@ -0,0 +1,47 @@
+import chromadb
+from sentence_transformers import SentenceTransformer
+import sqlite3
+
+def search(query: str, top_k: int = 5):
+    """
+    Searches the vector store for the most relevant documents to a given query.
+    """
+    # Handle special case for listing all students
+    if "give me the names of the students" in query.lower():
+        conn = sqlite3.connect('database/university.db')
+        cursor = conn.cursor()
+        students = cursor.execute("SELECT name FROM students").fetchall()
+        conn.close()
+        return [{"name": student[0]} for student in students]
+
+    # Initialize the embedding model
+    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+    # Create the query embedding
+    query_embedding = model.encode([query])
+
+    # Initialize ChromaDB client and get the collection
+    client = chromadb.PersistentClient(path="rag/vector_store")
+    collection = client.get_collection("university_data")
+
+    # Perform the search
+    results = collection.query(
+        query_embeddings=query_embedding,
+        n_results=top_k
+    )
+
+    return results['documents'][0]
+
+if __name__ == '__main__':
+    # Example usage
+    test_query = "What courses are available?"
+    results = search(test_query)
+    print(f"Results for '{test_query}':")
+    for result in results:
+        print(result)
+
+    test_query_students = "Give me the names of the students"
+    results_students = search(test_query_students)
+    print(f"\nResults for '{test_query_students}':")
+    for student in results_students:
+        print(student['name'])