token-optimise 0.1.0__py3-none-any.whl

src/core/cache.py

@@ -0,0 +1,126 @@
+import hashlib, time, os, logging
+from src.core.client import _chroma_client, CHROMA_PATH 
+from chromadb.utils import embedding_functions
+from token_optimise.config import settings, PROJECT_ROOT
+
+logger = logging.getLogger("token")
+
+
+TTL_BY_TOOL = {
+    "execute": 300,
+    "ask_document": 86400,
+    "search_all_documents": 86400,
+    "list_indexed_documents": 60,
+    "index_document": 0,   
+    "index_documents_folder": 0,
+}
+TTL_DEFAULT = 3600
+
+_cache_collection = None
+
+def _get_cache():
+    global _cache_collection
+    if _cache_collection is None:
+        embedder = embedding_functions.SentenceTransformerEmbeddingFunction(
+            model_name=settings.embedder
+        )
+        _cache_collection = _chroma_client.get_or_create_collection(
+            name="semantic_cache", embedding_function=embedder
+        )
+    return _cache_collection
+
+def check_cache(query: str, tool_name: str = ""):
+    """Returns (answer, similarity). answer is None on miss or expiry."""
+    try:
+        results = _get_cache().query(query_texts=[query], n_results=1)
+        if not results["documents"][0]:
+            return None, 0.0
+        distance = results["distances"][0][0]
+        similarity = 1 - distance
+        if similarity >= 0.8:
+            meta = results["metadatas"][0][0]
+            cached_at = meta.get("cached_at", 0)
+            ttl = TTL_BY_TOOL.get(tool_name, TTL_DEFAULT)
+            if ttl == 0:
+                return None, similarity
+            age = time.time() - cached_at
+            if age > ttl:
+                logger.info(
+                    f"[CACHE] EXPIRED for '{tool_name}' | age={int(age)}s ttl={ttl}s"
+                )
+                return None, similarity
+            logger.info(f"[CACHE] HIT for '{tool_name}' | sim={similarity:.3f} age={int(age)}s")
+            return meta["answer"], similarity
+        return None, similarity
+    except Exception as e:
+        logger.warning(f"[CACHE] check failed: {e}", exc_info=True)
+        return None, 0.0
+
+
+def check_cache_exact(key: str, tool_name: str = ""):
+    try:
+        result = _get_cache().get(ids=[key], include=["metadatas"])
+        if not result["ids"]:
+            return None, 0.0
+        meta = result["metadatas"][0]
+        cached_at = meta.get("cached_at", 0)
+        ttl = TTL_BY_TOOL.get(tool_name, TTL_DEFAULT)
+        if ttl == 0:
+            return None, 1.0
+        age = time.time() - cached_at
+        if age > ttl:
+            logger.info(f"[CACHE] EXPIRED (exact) for '{tool_name}' | age={int(age)}s")
+            return None, 1.0
+        logger.info(f"[CACHE] HIT (exact) for '{tool_name}'")
+        return meta["answer"], 1.0
+    except Exception as e:
+        logger.warning(f"[CACHE] exact check failed: {e}")
+        return None, 0.0
+    
+    
+    
+
+def store_answer(query, answer,tool_name):
+    query_id = hashlib.sha256(query.strip().lower().encode()).hexdigest()
+    try:
+        _get_cache().upsert(
+            ids=[query_id],
+            documents=[query],
+            metadatas=[{"answer": answer, "cached_at": time.time(), "tool_name": tool_name}],
+        )
+    except Exception as e:
+        logger.warning(f"[CACHE] store failed: {e}")
+        
+
+def store_answer_exact(key: str, answer: str, tool_name: str) -> None:
+    """Store by exact key (no embedding). Used for doc queries."""
+    try:
+        _get_cache().upsert(
+            ids=[key],
+            documents=[key],  
+            metadatas=[{"answer": answer, "cached_at": time.time(), "tool_name": tool_name}],
+        )
+    except Exception as e:
+        logger.warning(f"[CACHE] exact store failed: {e}")
+
+
+
+def cleanup_cache():
+    try:
+        all_entries = _get_cache().get(include=["metadatas"])
+        now = time.time()
+        expired_ids = []
+        for id_, meta in zip(all_entries["ids"], all_entries["metadatas"]):
+            tool_name = meta.get("tool_name", "")
+            ttl = TTL_BY_TOOL.get(tool_name, TTL_DEFAULT)
+            if ttl == 0:
+                continue
+            if now - meta.get("cached_at", 0) > ttl:
+                expired_ids.append(id_)
+        if expired_ids:
+            _get_cache().delete(ids=expired_ids)
+            logger.info(f"[CACHE] cleaned up {len(expired_ids)} expired entries")
+        return len(expired_ids)
+    except Exception as e:
+        logger.warning(f"[CACHE] cleanup failed: {e}")
+        return 0

src/core/client.py

@@ -0,0 +1,114 @@
+from groq import Groq
+from token_optimise.config import settings,PROJECT_ROOT
+import chromadb
+import os,json
+import logging
+
+logger = logging.getLogger("token")
+
+_groq_client = Groq(api_key=settings.groq_api_key)
+CHROMA_PATH = os.path.join(PROJECT_ROOT, "storage", "chroma_db")
+_chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
+
+def fill_args_llm(query,schema):
+    props = schema.get("properties", {})
+    required = schema.get("required", [])
+    schema_summary="\n".join([
+        f"- {name} ({info.get('type','string')}) : {info.get('description','no desc')}"
+        for name,info in props.items()
+    ])
+    
+    prompt = (
+        f"You are a tool argument filler. Given a tool schema and a user query, "
+        f"return ONLY a valid JSON object with the correct arguments.\n"
+        f"Tool parameters:\n{schema_summary}\n"
+        f"Required fields: {required}\n"
+        f"User query: \"{query}\"\n\n"
+        f"Rules:\n"
+        f"- Return ONLY a JSON object, no explanation, no markdown, no backticks\n"
+        f"- For file paths: NEVER use ~ or relative paths. Always expand to full absolute path.\n"
+        f"- Home directory is: {os.path.expanduser('~')}\n"
+        f"- Desktop is: {os.path.expanduser('~/Desktop')}\n"
+        f"- Documents is: {os.path.expanduser('~/Documents')}\n"
+        f"JSON:"
+    )
+        
+    
+    try:
+        raw = _groq_client.chat.completions.create(
+            model=settings.groq_model,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0,
+            max_tokens=settings.max_response_tokens,
+        )
+        usage = raw.usage
+        logger.info(
+            f"[GROQ] fill_args | prompt_tokens={usage.prompt_tokens} "
+            f"completion_tokens={usage.completion_tokens}"
+        )
+
+        response = raw.choices[0].message.content.strip()
+        
+        if response.startswith("```"):
+            response = response.split("```")[1]
+            if response.startswith("json"):
+                response = response[4:]
+            response = response.strip()
+
+        args = json.loads(response)
+        result = {k: v for k, v in args.items() if v is not None}
+        result["_groq_usage"] = {
+            "prompt_tokens": usage.prompt_tokens,
+            "completion_tokens": usage.completion_tokens,
+        }
+        return result
+
+    except json.JSONDecodeError as e:
+        logger.warning(f"[GROQ] failed to parse args JSON: {e}")
+        return {"_groq_error": f"JSONDecodeError: {e}"}
+    except Exception as e:
+        logger.warning(f"[GROQ] fill_args_llm failed: {e}")
+        return {"_groq_error": str(e)}
+
+
+def expand_query(query: str, tool_descriptions: list[dict]) -> tuple[str, int, int]:
+    """Rewrite query using tool terminology. Returns (rewritten, prompt_tokens, completion_tokens)."""
+    if not tool_descriptions:
+        return query, 0, 0
+
+    # Using all the tools, capping stopped.
+    tools_text = "\n".join(
+        f"- {t['function']['name']}: {t['function'].get('description', '')}"
+        for t in tool_descriptions
+    )
+
+    try:
+        resp = _groq_client.chat.completions.create(
+            model=settings.groq_model,
+            messages=[
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a query rewriter. Given a user query and a list of available tools, "
+                        "rewrite the query using the exact terminology and phrasing that best matches "
+                        "the tool descriptions. Output ONLY the rewritten query, nothing else."
+                    ),
+                },
+                {
+                    "role": "user",
+                    "content": f"Tools:\n{tools_text}\n\nUser query: {query}\n\nRewritten query:",
+                },
+            ],
+            max_tokens=60,
+            temperature=0.0,
+        )
+        usage = resp.usage
+        logger.info(
+            f"[GROQ] expand_query | prompt_tokens={usage.prompt_tokens} "
+            f"completion_tokens={usage.completion_tokens}"
+        )
+        rewritten = resp.choices[0].message.content.strip()
+        return (rewritten if rewritten else query, usage.prompt_tokens, usage.completion_tokens)
+    except Exception as e:
+        logger.warning(f"[EXPAND_QUERY] failed: {e}")
+        return query, 0, 0

src/core/db.py

@@ -0,0 +1,299 @@
+import os
+import time
+import uuid
+import logging
+import threading
+import sqlite3
+import requests
+from contextlib import contextmanager
+from token_optimise.config import PROJECT_ROOT,settings
+
+logger = logging.getLogger("token")
+
+DB_PATH = os.path.join(PROJECT_ROOT, "storage", "token_events.db")
+
+INGEST_URL = settings.token_ingest_url
+
+
+CONVERSATION_TIMEOUT_MINUTES = 15
+_current_conversation_id = None
+_last_event_time = None
+_conv_lock = threading.Lock()
+
+
+def _get_or_create_conversation_id() -> str:
+    global _current_conversation_id, _last_event_time
+    with _conv_lock:
+        now = time.time()
+        if (
+            _current_conversation_id is None
+            or _last_event_time is None
+            or (now - _last_event_time) > CONVERSATION_TIMEOUT_MINUTES * 60
+        ):
+            _current_conversation_id = str(uuid.uuid4())
+        _last_event_time = now
+        return _current_conversation_id
+
+
+def get_connection() -> sqlite3.Connection:
+    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+@contextmanager
+def get_db():
+    conn = get_connection()
+    try:
+        yield conn
+    finally:
+        conn.close()
+
+
+def init_db(conn):
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS events (
+            id  INTEGER PRIMARY KEY AUTOINCREMENT,
+            user_id     TEXT,
+            timestamp   DATETIME DEFAULT CURRENT_TIMESTAMP,
+            conversation_id     TEXT,
+            tool_name   TEXT NOT NULL,
+            query   TEXT,
+            cache_hit   INTEGER DEFAULT 0,
+            cache_similarity    REAL DEFAULT 0.0,
+            tokens_before_trim  INTEGER DEFAULT 0,
+            tokens_after_trim   INTEGER DEFAULT 0,
+            trim_saved  INTEGER DEFAULT 0,
+            schema_tokens_full   INTEGER DEFAULT 0,
+            schema_tokens_selected INTEGER DEFAULT 0,
+            schema_tokens_saved INTEGER DEFAULT 0,
+            groq_prompt_tokens INTEGER DEFAULT 0,
+            groq_completion_tokens  INTEGER DEFAULT 0,
+            doc_id  TEXT,
+            success INTEGER DEFAULT 1
+        )
+    """)
+    for col in ["user_id", "conversation_id"]:
+        try:
+            conn.execute(f"ALTER TABLE events ADD COLUMN {col} TEXT")
+        except Exception:
+            pass
+    conn.commit()
+    logger.info(f"[DB] SQLite ready at {DB_PATH}")
+
+
+
+def init_db_once() -> None:
+    """Call once at server startup."""
+    with get_db() as conn:
+        init_db(conn)
+
+
+
+
+
+def _post_to_cloud(data: dict):
+    """Send event to AWS /collect endpoint."""
+    if not INGEST_URL:
+        return
+    try:
+        # Read token fresh — it may have been saved after module loaded
+        token = settings.collect_token or ""
+        headers = {"X-Collect-Token": token} if token else {}
+        requests.post(INGEST_URL, json=data, timeout=3, headers=headers)
+    except Exception as e:
+        logger.debug(f"[DB] Cloud ingest failed (local SQLite has it): {e}")
+
+
+
+def insert_event(**kwargs):
+    kwargs.setdefault("user_id", os.getenv("TOKEN_USER_ID", "unknown"))
+    kwargs.setdefault("conversation_id", _get_or_create_conversation_id())
+
+    fields = [
+        "user_id", "conversation_id",
+        "tool_name", "query", "cache_hit", "cache_similarity",
+        "tokens_before_trim", "tokens_after_trim", "trim_saved",
+        "schema_tokens_full", "schema_tokens_selected", "schema_tokens_saved",
+        "doc_id", "success",
+        "groq_prompt_tokens", "groq_completion_tokens",
+    ]
+
+    data = {f: kwargs.get(f, None) for f in fields}
+
+    # always write to local SQLite
+    with get_db() as conn:
+        placeholders = ", ".join(["?" for _ in fields])
+        columns = ", ".join(fields)
+        conn.execute(
+            f"INSERT INTO events ({columns}) VALUES ({placeholders})",
+            list(data.values()),
+        )
+        conn.commit()
+
+
+    _post_to_cloud(data)
+        
+        
+def _fetchall(conn, sql, params=()):
+    rows = conn.execute(sql, params).fetchall()
+    return [dict(r) for r in rows]
+
+
+
+def _fetchone(conn, sql, params=()):
+    row = conn.execute(sql, params).fetchone()
+    return dict(row) if row else {}
+
+         
+def get_summary(conn):
+    return _fetchone(conn, """
+        SELECT
+            COUNT(*) as total_calls,
+            COALESCE(SUM(cache_hit), 0) as cache_hits,
+            COALESCE(ROUND(AVG(cache_hit) * 100, 1), 0.0) as hit_rate_pct,
+            COALESCE(SUM(trim_saved), 0) as total_trim_saved,
+            COALESCE(SUM(schema_tokens_saved), 0) as total_schema_saved,
+            COALESCE(SUM(trim_saved + schema_tokens_saved), 0) as total_tokens_saved
+        FROM events
+    """)
+
+
+
+def get_recent_events(conn, limit=50):
+    return _fetchall(
+        conn, "SELECT * FROM events ORDER BY timestamp DESC LIMIT ?", (limit,)
+    )
+
+def get_tool_stats(conn):
+    return _fetchall(conn,"""
+        SELECT 
+            tool_name,
+            COUNT(*) AS calls,
+            SUM(cache_hit) AS hits,
+            SUM(trim_saved) AS trim_saved,
+            SUM(schema_tokens_saved) AS schema_saved
+        FROM events
+        GROUP BY tool_name
+        ORDER BY calls DESC
+    """)
+    
+
+def get_token_analysis(conn):
+    rows = _fetchone(conn, """
+        SELECT 
+            COUNT(*) AS total_queries,
+            COALESCE(SUM(schema_tokens_saved), 0) AS schema_saved,
+            COALESCE(SUM(schema_tokens_full), 0) AS schema_full,
+            COALESCE(SUM(schema_tokens_selected), 0) AS schema_selected,
+            COALESCE(SUM(trim_saved), 0) AS trim_saved,
+            COALESCE(SUM(tokens_before_trim), 0) AS tokens_before_trim,
+            COALESCE(SUM(tokens_after_trim), 0) AS tokens_after_trim,
+            COALESCE(SUM(cache_hit), 0) AS cache_hits,
+            COUNT(CASE WHEN cache_hit = 0 THEN 1 END) as cache_misses
+        FROM events
+    """)
+    
+    schema_saved = rows.get("schema_saved",0)
+    schema_full = rows.get("schema_full",0)
+    schema_selected = rows.get("schema_selected",0)
+    trim_saved = rows.get("trim_saved",0)
+    tokens_before = rows.get("tokens_before_trim",0)
+    tokens_after = rows.get("tokens_after_trim",0)
+    cache_hits = rows.get("cache_hits",0)
+    cache_misses = rows.get("cache_misses",0)
+    total_queries = rows.get("total_queries",0)
+    
+    total_saved = schema_saved + trim_saved
+    
+    actual_without= schema_full + tokens_before
+    actual_with = schema_selected + tokens_after
+    
+    pct_saved = round((total_saved / actual_without * 100), 1) if actual_without > 0 else 0.0
+
+    return {
+        "total_queries": total_queries,
+        "cache_hits": cache_hits,
+        "cache_misses": cache_misses,
+        "schema_tokens_without_tom": schema_full,
+        "schema_tokens_with_tom": schema_selected,
+        "schema_saved": schema_saved,
+        "tokens_before_trim": tokens_before,
+        "tokens_after_trim": tokens_after,
+        "trim_saved": trim_saved,
+        "total_saved": total_saved,
+        "actual_without_tom": actual_without,
+        "actual_with_tom": actual_with,
+        "pct_saved": pct_saved,
+    }
+    
+def get_event_count(conn) -> int:
+    """Quick row count for health check."""
+    row = _fetchone(conn, "SELECT COUNT(*) AS cnt FROM events")
+    return row.get("cnt", 0)
+
+def get_conversation_stats(conn, limit=20):
+    return _fetchall(conn, """
+        SELECT
+            conversation_id,
+            MIN(timestamp) AS started_at,
+            COUNT(*) AS total_calls,
+            SUM(cache_hit)  AS cache_hits,
+            SUM(trim_saved) AS trim_saved,
+            SUM(schema_tokens_saved) AS schema_saved,
+            SUM(trim_saved + schema_tokens_saved) AS total_saved
+        FROM events
+        WHERE conversation_id IS NOT NULL
+        GROUP BY conversation_id
+        ORDER BY started_at DESC
+        LIMIT ?
+    """, (limit,))
+
+    
+    
+def get_groq_usage(conn):
+    return _fetchone(conn, """
+        SELECT
+            COALESCE(SUM(groq_prompt_tokens), 0) AS total_prompt_tokens,
+            COALESCE(SUM(groq_completion_tokens), 0)  AS total_completion_tokens,
+            COALESCE(SUM(groq_prompt_tokens + groq_completion_tokens), 0) AS total_groq_tokens
+        FROM events
+    """)
+
+
+def get_timeseries(conn, hours: int = 24):
+    return _fetchall(conn, """
+        SELECT
+            strftime('%Y-%m-%d %H:00', timestamp) AS hour,
+            COUNT(*) AS total_calls,
+            COALESCE(SUM(cache_hit), 0)  AS cache_hits,
+            COALESCE(SUM(schema_tokens_saved), 0) AS schema_saved,
+            COALESCE(SUM(trim_saved), 0) AS trim_saved,
+            COALESCE(SUM(groq_prompt_tokens + groq_completion_tokens), 0) AS groq_tokens
+        FROM events
+        WHERE timestamp >= datetime('now', ? || ' hours')
+        GROUP BY hour
+        ORDER BY hour ASC
+    """, (f"-{hours}",))
+    
+
+    
+def get_period_summary(conn, period: str = "today"):
+    _safe = {
+        "today": "date(timestamp) = date('now')",
+        "yesterday": "date(timestamp) = date('now', '-1 day')",
+        "week": "timestamp >= datetime('now', '-7 days')",
+        "all": "1=1",
+    }
+    where = _safe.get(period, "1=1")
+    return _fetchone(conn, f"""
+        SELECT
+            COUNT(*) AS total_calls,
+            COALESCE(SUM(cache_hit), 0) AS cache_hits,
+            COALESCE(SUM(schema_tokens_saved + trim_saved), 0) AS total_saved,
+            COALESCE(SUM(groq_prompt_tokens), 0)  AS groq_prompt,
+            COALESCE(SUM(groq_completion_tokens), 0) AS groq_completion
+        FROM events
+        WHERE {where}
+    """)