diff --git a/chatbot/response_generator.py b/chatbot/response_generator.py index 85c4dfc..74cbac2 100644 --- a/chatbot/response_generator.py +++ b/chatbot/response_generator.py @@ -2,10 +2,17 @@ from __future__ import annotations +import re import requests from typing import Iterable, Mapping from ai.rag_pipeline import RAGPipeline +from db.query_engine import exact_lookup + + +def looks_like_exact_lookup(message: str) -> bool: + """Return True if the message contains an ISO date-time pattern.""" + return bool(re.search(r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}", message)) class OllamaModel: @@ -38,6 +45,25 @@ def __init__(self, tenant_id: str, model_type: str = "ollama", model_name: str = else: raise ValueError(f"Unsupported model type: {model_type}") - def generate_response(self, user_message: str, history: Iterable[Mapping[str, str]] | None = None) -> str: + def query_rag(self, user_message: str, history: Iterable[Mapping[str, str]] | None = None) -> str: + """Fallback to the RAG pipeline and language model.""" prompt = self.rag.augment_prompt(user_message, history) return self.model.generate(prompt) + + def generate_response(self, user_message: str, history: Iterable[Mapping[str, str]] | None = None) -> str: + """Generate a response handling exact lookup or RAG fallback.""" + if looks_like_exact_lookup(user_message): + date_match = re.search(r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}", user_message) + if date_match: + date_str = date_match.group(0) + row = exact_lookup(date_str, self.tenant_id) + if row: + return ( + f"Close value for {date_str} is {row['close']} " + f"(Open: {row['open']}, High: {row['high']}, " + f"Low: {row['low']}, Volume: {row['volume']})" + ) + else: + return f"No data found for {date_str}." + + return self.query_rag(user_message, history) diff --git a/data/tenant2.db b/data/tenant2.db new file mode 100644 index 0000000..e69de29 diff --git a/data/testtenant.db b/data/testtenant.db new file mode 100644 index 0000000..fb0c9a5 Binary files /dev/null and b/data/testtenant.db differ diff --git a/db/query_engine.py b/db/query_engine.py index 2fc703d..deaa74d 100644 --- a/db/query_engine.py +++ b/db/query_engine.py @@ -36,3 +36,34 @@ def execute(self, query: str, params=None): def close(self) -> None: """Close the active connector.""" self.connector.close() + + +def exact_lookup(date_str: str, tenant_id: str, table: str = "market_data") -> dict: + """Retrieve exact row from the tenant's database by date. + + Args: + date_str: Date-time string in ISO format "YYYY-MM-DD HH:MM" + tenant_id: Current tenant identifier + table: Target table name + + Returns: + dict: Matching row or {} if not found + """ + import sqlite3 + + db_path = f"data/{tenant_id}.db" + try: + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + query = f"SELECT * FROM {table} WHERE date = ? LIMIT 1" + cursor.execute(query, (date_str,)) + row = cursor.fetchone() + conn.close() + except Exception: + return {} + + if not row: + return {} + + columns = ["date", "open", "high", "low", "close", "volume"] + return dict(zip(columns, row)) diff --git a/scripts/check_db.py b/scripts/check_db.py new file mode 100644 index 0000000..9fbab15 --- /dev/null +++ b/scripts/check_db.py @@ -0,0 +1,25 @@ +"""Simple helper to check a tenant's market_data table for a specific date.""" + +import sqlite3 +import sys + + +def main(tenant_id: str, date_to_check: str) -> None: + db_path = f"data/{tenant_id}.db" + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute("SELECT * FROM market_data WHERE date = ?", (date_to_check,)) + row = cursor.fetchone() + conn.close() + if row: + print(row) + else: + print("No data found") + + +if __name__ == "__main__": + if len(sys.argv) != 3: + print("Usage: python scripts/check_db.py ") + sys.exit(1) + main(sys.argv[1], sys.argv[2]) + diff --git a/scripts/load_tenant_data.py b/scripts/load_tenant_data.py index 5908f84..a1beac3 100644 --- a/scripts/load_tenant_data.py +++ b/scripts/load_tenant_data.py @@ -1,6 +1,8 @@ -"""Utility script to seed tenant2's vector store with example data.""" +"""Utility script to seed tenant2's data stores with example data.""" from pathlib import Path +import os +import sqlite3 import sys # Ensure the project root is on the Python path so ``ai`` imports work when @@ -9,13 +11,39 @@ from ai.vector_stores.chroma_store import TenantVectorStore -def load_tenant2_data(): + +def load_tenant2_data() -> None: + """Populate tenant2's vector store and SQLite DB with sample data.""" + store = TenantVectorStore("tenant2") - store.add_document("doc1", "Dan middle name is the king69$$$, he is 186 cm tall and he likes banana flavoured ice cream") - store.add_document("doc2", "Bitcoin is down past 115000 usd, showing that distribution started last week around july 25th") + store.add_document( + "doc1", + "Dan middle name is the king69$$$, he is 186 cm tall and he likes banana flavoured ice cream", + ) + store.add_document( + "doc2", + "Bitcoin is down past 115000 usd, showing that distribution started last week around july 25th", + ) print("\u2705 Data loaded for tenant2") + os.makedirs("data", exist_ok=True) + db_path = "data/tenant2.db" + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE IF NOT EXISTS market_data (date TEXT, open REAL, high REAL, low REAL, close REAL, volume REAL)" + ) + cursor.execute( + "INSERT OR IGNORE INTO market_data VALUES (?, ?, ?, ?, ?, ?)", + ("2023-07-09 22:01", 10.0, 12.0, 9.0, 11.0, 1000.0), + ) + conn.commit() + conn.close() + print(f"\u2705 market_data table initialized at {db_path}") + + if __name__ == "__main__": load_tenant2_data() + diff --git a/tests/test_exact_lookup.py b/tests/test_exact_lookup.py new file mode 100644 index 0000000..17c99ad --- /dev/null +++ b/tests/test_exact_lookup.py @@ -0,0 +1,49 @@ +import os +import sqlite3 + +from chatbot.response_generator import ResponseGenerator +from db.query_engine import exact_lookup + + +TENANT = "testtenant" +DATE_STR = "2023-07-09 22:01" + + +def setup_module(module): + """Create a temporary SQLite DB for the tenant.""" + os.makedirs("data", exist_ok=True) + db_path = f"data/{TENANT}.db" + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE IF NOT EXISTS market_data (date TEXT, open REAL, high REAL, low REAL, close REAL, volume REAL)" + ) + cursor.execute( + "INSERT INTO market_data VALUES (?, ?, ?, ?, ?, ?)", + (DATE_STR, 10.0, 12.0, 9.0, 11.0, 1000.0), + ) + conn.commit() + conn.close() + + +def test_exact_lookup_function(): + row = exact_lookup(DATE_STR, TENANT) + assert row["close"] == 11.0 + + +def test_response_generator_exact_and_rag(): + rg = ResponseGenerator(TENANT) + + def fake_query_rag(self, user_message, history=None): + return "RAG" + + # Replace query_rag to avoid actual model calls + rg.query_rag = fake_query_rag.__get__(rg, ResponseGenerator) + + exact_msg = f"What is the close value for date: {DATE_STR}?" + response = rg.generate_response(exact_msg) + assert "Close value for" in response + + general_msg = "What happened on 2023-07-09?" + assert rg.generate_response(general_msg) == "RAG" + diff --git a/vector_store/ctx/06b71e95-9239-4ec6-8213-200a4afec9d0/length.bin b/vector_store/ctx/06b71e95-9239-4ec6-8213-200a4afec9d0/length.bin index de41e66..c2bd53b 100644 Binary files a/vector_store/ctx/06b71e95-9239-4ec6-8213-200a4afec9d0/length.bin and b/vector_store/ctx/06b71e95-9239-4ec6-8213-200a4afec9d0/length.bin differ diff --git a/vector_store/ctx/chroma.sqlite3 b/vector_store/ctx/chroma.sqlite3 index b51e3c0..cb813ed 100644 Binary files a/vector_store/ctx/chroma.sqlite3 and b/vector_store/ctx/chroma.sqlite3 differ diff --git a/vector_store/ol/chroma.sqlite3 b/vector_store/ol/chroma.sqlite3 index fa09648..aa796a2 100644 Binary files a/vector_store/ol/chroma.sqlite3 and b/vector_store/ol/chroma.sqlite3 differ diff --git a/vector_store/rag/1b70bdf9-03f9-42ad-a18e-bcb6e9ae72dc/length.bin b/vector_store/rag/1b70bdf9-03f9-42ad-a18e-bcb6e9ae72dc/length.bin index 26a7ca7..b5ec4d9 100644 Binary files a/vector_store/rag/1b70bdf9-03f9-42ad-a18e-bcb6e9ae72dc/length.bin and b/vector_store/rag/1b70bdf9-03f9-42ad-a18e-bcb6e9ae72dc/length.bin differ diff --git a/vector_store/rag/chroma.sqlite3 b/vector_store/rag/chroma.sqlite3 index 536d688..90c8cca 100644 Binary files a/vector_store/rag/chroma.sqlite3 and b/vector_store/rag/chroma.sqlite3 differ diff --git a/vector_store/t1/56dc3db9-3bb6-4e63-9e3f-02e0fce17e67/length.bin b/vector_store/t1/56dc3db9-3bb6-4e63-9e3f-02e0fce17e67/length.bin index cd8ed3b..96e02fe 100644 Binary files a/vector_store/t1/56dc3db9-3bb6-4e63-9e3f-02e0fce17e67/length.bin and b/vector_store/t1/56dc3db9-3bb6-4e63-9e3f-02e0fce17e67/length.bin differ diff --git a/vector_store/t1/chroma.sqlite3 b/vector_store/t1/chroma.sqlite3 index e02f924..dce6389 100644 Binary files a/vector_store/t1/chroma.sqlite3 and b/vector_store/t1/chroma.sqlite3 differ diff --git a/vector_store/tenant2/chroma.sqlite3 b/vector_store/tenant2/chroma.sqlite3 index 59cc0ef..27ecce8 100644 Binary files a/vector_store/tenant2/chroma.sqlite3 and b/vector_store/tenant2/chroma.sqlite3 differ diff --git a/vector_store/testtenant/chroma.sqlite3 b/vector_store/testtenant/chroma.sqlite3 new file mode 100644 index 0000000..7ff3d8f Binary files /dev/null and b/vector_store/testtenant/chroma.sqlite3 differ