Add exact DB lookup for date-based queries with RAG fallback

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)

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))

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 <tenant_id> <date>")
+        sys.exit(1)
+    main(sys.argv[1], sys.argv[2])
+

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()
+

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"
+