cli-memory-os 0.1.0__py3-none-any.whl

storage/db.py

@@ -0,0 +1,528 @@
+import sqlite3
+import os
+import logging
+DB_PATH = "memory.db"
+SCHEMA_PATH = os.path.join(os.path.dirname(__file__), "schema.sql")
+
+_initialized_dbs = set()
+
+def get_connection():
+    global _initialized_dbs
+    db_path = os.getenv("MEMORY_OS_DB_PATH", DB_PATH)
+    abs_path = os.path.abspath(db_path)
+    
+    # Ensure directory exists
+    parent = os.path.dirname(abs_path)
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+    
+    conn = sqlite3.connect(abs_path)
+    conn.row_factory = sqlite3.Row
+    
+    if abs_path not in _initialized_dbs:
+        # Check if the schema is initialized (e.g. check if document_chunks table exists)
+        cursor = conn.cursor()
+        try:
+            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='document_chunks'")
+            table_exists = cursor.fetchone()
+        except sqlite3.OperationalError:
+            table_exists = False
+            
+        if not table_exists:
+            _initialized_dbs.add(abs_path)
+            conn.close()
+            init_db()
+            conn = sqlite3.connect(abs_path)
+            conn.row_factory = sqlite3.Row
+        else:
+            _initialized_dbs.add(abs_path)
+            
+    return conn
+
+
+
+def init_db():
+    conn = get_connection()
+    cursor = conn.cursor()
+    # Check if repositories table exists and needs recreation (e.g. legacy schema)
+    cursor.execute("PRAGMA table_info(repositories)")
+    columns = [row[1] for row in cursor.fetchall()]
+    if columns and "visibility" not in columns:
+        cursor.execute("DROP TABLE repositories")
+        cursor.execute("DROP TABLE IF EXISTS repository_documents")
+        
+    # Check if emails table exists and needs recreation (e.g. legacy schema without message_id)
+    cursor.execute("PRAGMA table_info(emails)")
+    email_columns = [row[1] for row in cursor.fetchall()]
+    if email_columns and "message_id" not in email_columns:
+        cursor.execute("DROP TABLE emails")
+        
+    if os.path.exists(SCHEMA_PATH):
+        with open(SCHEMA_PATH, "r", encoding="utf-8") as f:
+            schema_sql = f.read()
+        conn.executescript(schema_sql)
+    conn.commit()
+    conn.close()
+
+def insert_repository(repo):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT id FROM repositories WHERE repo_name = ?", (repo.repo_name,))
+    row = cursor.fetchone()
+    if row:
+        cursor.execute(
+            """
+            UPDATE repositories 
+            SET description = ?, language = ?, visibility = ?, stars = ?, forks = ?, open_issues = ?, default_branch = ?, updated_at = ?, url = ?
+            WHERE id = ?
+            """,
+            (repo.description, repo.language, repo.visibility, repo.stars, repo.forks, repo.open_issues, repo.default_branch, repo.updated_at, repo.url, row[0])
+        )
+    else:
+        cursor.execute(
+            """
+            INSERT INTO repositories (repo_name, description, language, visibility, stars, forks, open_issues, default_branch, updated_at, url)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (repo.repo_name, repo.description, repo.language, repo.visibility, repo.stars, repo.forks, repo.open_issues, repo.default_branch, repo.updated_at, repo.url)
+        )
+    conn.commit()
+    conn.close()
+
+def insert_repository_document(doc):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        "SELECT id FROM repository_documents WHERE repo_name = ? AND file_name = ?",
+        (doc.repo_name, doc.file_name)
+    )
+    row = cursor.fetchone()
+    if row:
+        cursor.execute(
+            "UPDATE repository_documents SET content = ?, source = ?, synced_at = ? WHERE id = ?",
+            (doc.content, doc.source, doc.synced_at, row[0])
+        )
+    else:
+        cursor.execute(
+            "INSERT INTO repository_documents (repo_name, file_name, content, source, synced_at) VALUES (?, ?, ?, ?, ?)",
+            (doc.repo_name, doc.file_name, doc.content, doc.source, doc.synced_at)
+        )
+    conn.commit()
+    conn.close()
+
+def insert_email(email):
+    conn = get_connection()
+    cursor = conn.cursor()
+    if email.message_id:
+        cursor.execute("SELECT id FROM emails WHERE message_id = ?", (email.message_id,))
+    else:
+        cursor.execute(
+            "SELECT id FROM emails WHERE subject = ? AND sender = ? AND received_at = ?",
+            (email.subject, email.sender, email.received_at)
+        )
+    row = cursor.fetchone()
+    if row:
+        cursor.execute(
+            "UPDATE emails SET snippet = ?, subject = ?, sender = ?, received_at = ? WHERE id = ?",
+            (email.snippet, email.subject, email.sender, email.received_at, row[0])
+        )
+    else:
+        cursor.execute(
+            "INSERT INTO emails (message_id, subject, sender, snippet, received_at) VALUES (?, ?, ?, ?, ?)",
+            (email.message_id, email.subject, email.sender, email.snippet, email.received_at)
+        )
+    conn.commit()
+    conn.close()
+
+def get_repo_count() -> int:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT COUNT(*) FROM repositories")
+    count = cursor.fetchone()[0]
+    conn.close()
+    return count
+
+def get_repository_document_count() -> int:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT COUNT(*) FROM repository_documents")
+    count = cursor.fetchone()[0]
+    conn.close()
+    return count
+
+def get_email_count() -> int:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT COUNT(*) FROM emails")
+    count = cursor.fetchone()[0]
+    conn.close()
+    return count
+
+def get_repository_details(repo_name: str):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        "SELECT repo_name, description, language, visibility, stars, forks, open_issues, default_branch, updated_at, url FROM repositories WHERE LOWER(repo_name) = LOWER(?)",
+        (repo_name,)
+    )
+    repo_row = cursor.fetchone()
+    if not repo_row:
+        conn.close()
+        return None
+    
+    matched_name = repo_row[0]
+    
+    # Get files stored for this repository
+    cursor.execute(
+        "SELECT file_name FROM repository_documents WHERE LOWER(repo_name) = LOWER(?)",
+        (repo_name,)
+    )
+    files = [row[0] for row in cursor.fetchall()]
+    
+    # Get README content
+    cursor.execute(
+        "SELECT content FROM repository_documents WHERE LOWER(repo_name) = LOWER(?) AND LOWER(file_name) = 'readme.md'",
+        (repo_name,)
+    )
+    readme_row = cursor.fetchone()
+    readme_content = readme_row[0] if readme_row else None
+    
+    conn.close()
+    return {
+        "repo_name": matched_name,
+        "description": repo_row[1],
+        "language": repo_row[2],
+        "visibility": repo_row[3],
+        "stars": repo_row[4],
+        "forks": repo_row[5],
+        "open_issues": repo_row[6],
+        "default_branch": repo_row[7],
+        "updated_at": repo_row[8],
+        "url": repo_row[9],
+        "files": files,
+        "readme": readme_content
+    }
+
+def clear_all():
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("DELETE FROM repositories")
+    cursor.execute("DELETE FROM repository_documents")
+    cursor.execute("DELETE FROM emails")
+    cursor.execute("DELETE FROM document_chunks")
+    cursor.execute("DELETE FROM graph_relationships")
+    cursor.execute("DELETE FROM graph_nodes")
+    conn.commit()
+    conn.close()
+
+def search_local_knowledge_ranked(query: str, repo_filter: str = None) -> list:
+    """Ranked search across repositories, documents, and emails.
+
+    Applies configurable boost for repository scores and weight for email scores.
+    Adds ``source_type`` to each result and logs detailed diagnostics.
+    """
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    if repo_filter:
+        if isinstance(repo_filter, list):
+            placeholders = ",".join(["?"] * len(repo_filter))
+            cursor.execute(
+                f"SELECT repo_name, language, description FROM repositories WHERE LOWER(repo_name) IN ({placeholders})",
+                [r.lower() for r in repo_filter]
+            )
+            repos = cursor.fetchall()
+
+            cursor.execute(
+                f"SELECT repo_name, file_name, content FROM repository_documents WHERE LOWER(repo_name) IN ({placeholders})",
+                [r.lower() for r in repo_filter]
+            )
+            docs = cursor.fetchall()
+            
+            emails = []
+        else:
+            cursor.execute("SELECT repo_name, language, description FROM repositories WHERE LOWER(repo_name) = LOWER(?)", (repo_filter,))
+            repos = cursor.fetchall()
+
+            cursor.execute("SELECT repo_name, file_name, content FROM repository_documents WHERE LOWER(repo_name) = LOWER(?)", (repo_filter,))
+            docs = cursor.fetchall()
+
+            emails = []
+    else:
+        # Fetch all data for scoring
+        cursor.execute("SELECT repo_name, language, description FROM repositories")
+        repos = cursor.fetchall()
+
+        cursor.execute("SELECT repo_name, file_name, content FROM repository_documents")
+        docs = cursor.fetchall()
+
+        cursor.execute("SELECT subject, sender, snippet FROM emails")
+        emails = cursor.fetchall()
+    conn.close()
+
+    # Configurable boosts (default = 1.0 i.e., no change)
+    repo_boost = float(os.getenv('REPO_SCORE_BOOST', '1.0'))
+    email_weight = float(os.getenv('EMAIL_SCORE_WEIGHT', '1.0'))
+
+    ranked_results = []
+    query_lower = query.lower()
+    logger = logging.getLogger(__name__)
+
+    # --- Score repositories ---
+    for repo_name, language, description in repos:
+        score = 0
+        if repo_name and query_lower in repo_name.lower():
+            score += 10
+        if description and query_lower in description.lower():
+            score += 8
+        if language and query_lower in language.lower():
+            score += 5
+        if score > 0:
+            score = int(score * repo_boost)
+            result = {
+                "type": "repository",
+                "score": score,
+                "repo_name": repo_name,
+                "language": language,
+                "description": description,
+                "source_type": "repository",
+            }
+            ranked_results.append(result)
+            logger.debug(f"Retrieval result - {result}")
+
+    # --- Score documents ---
+    for repo_name, file_name, content in docs:
+        score = 0
+        content_lower = content.lower() if content else ""
+        file_name_lower = file_name.lower() if file_name else ""
+        if query_lower in file_name_lower or query_lower in content_lower:
+            if file_name_lower == "readme.md":
+                score += 6
+            elif file_name_lower == "package.json":
+                score += 4
+            else:
+                score += 3
+        if score > 0:
+            result = {
+                "type": "document",
+                "score": score,
+                "repo_name": repo_name,
+                "file_name": file_name,
+                "content": content,
+                "source_type": "document",
+            }
+            ranked_results.append(result)
+            logger.debug(f"Retrieval result - {result}")
+
+    # --- Score emails ---
+    for subject, sender, snippet in emails:
+        score = 0
+        subj_l = subject.lower() if subject else ""
+        send_l = sender.lower() if sender else ""
+        snip_l = snippet.lower() if snippet else ""
+        if query_lower in subj_l or query_lower in send_l or query_lower in snip_l:
+            score += 2
+        if score > 0:
+            score = int(score * email_weight)
+            result = {
+                "type": "email",
+                "score": score,
+                "subject": subject,
+                "sender": sender,
+                "snippet": snippet,
+                "source_type": "email",
+            }
+            ranked_results.append(result)
+            logger.debug(f"Retrieval result - {result}")
+
+    # Sort descending by score, then stable alphabetical tie‑breaker
+    ranked_results.sort(key=lambda x: (-x["score"], x.get("repo_name") or x.get("subject") or ""))
+    return ranked_results
+
+def get_repository_files(repo_name: str) -> list:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT file_name FROM repository_documents WHERE LOWER(repo_name) = LOWER(?)", (repo_name,))
+    files = [row[0] for row in cursor.fetchall()]
+    conn.close()
+    return files
+
+def get_repository_readme(repo_name: str) -> str:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT content FROM repository_documents WHERE LOWER(repo_name) = LOWER(?) AND LOWER(file_name) = 'readme.md'", (repo_name,))
+    row = cursor.fetchone()
+    readme = row[0] if row else None
+    conn.close()
+    return readme
+
+def get_all_repositories() -> list:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT repo_name, language, description, stars, forks, updated_at FROM repositories")
+    repos = [
+        {
+            "repo_name": row[0],
+            "language": row[1],
+            "description": row[2],
+            "stars": row[3],
+            "forks": row[4],
+            "updated_at": row[5]
+        }
+        for row in cursor.fetchall()
+    ]
+    conn.close()
+    return repos
+
+def get_all_documents() -> list:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT repo_name, file_name, content FROM repository_documents")
+    docs = [
+        {"repo_name": row[0], "file_name": row[1], "content": row[2]}
+        for row in cursor.fetchall()
+    ]
+    conn.close()
+    return docs
+
+def insert_document_chunk(repository_name: str, document_name: str, source_type: str, chunk_text: str, chunk_index: int, created_at: str):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        """
+        INSERT INTO document_chunks (repository_name, document_name, source_type, chunk_text, chunk_index, created_at)
+        VALUES (?, ?, ?, ?, ?, ?)
+        """,
+        (repository_name, document_name, source_type, chunk_text, chunk_index, created_at)
+    )
+    conn.commit()
+    conn.close()
+
+def clear_document_chunks():
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("DELETE FROM document_chunks")
+    cursor.execute("DELETE FROM sqlite_sequence WHERE name = 'document_chunks'")
+    conn.commit()
+    conn.close()
+
+def get_document_chunk_count() -> int:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT COUNT(*) FROM document_chunks")
+    count = cursor.fetchone()[0]
+    conn.close()
+    return count
+
+def get_all_document_chunks() -> list:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT id, repository_name, document_name, source_type, chunk_text, chunk_index FROM document_chunks")
+    chunks = [
+        {
+            "id": row[0],
+            "repository_name": row[1],
+            "document_name": row[2],
+            "source_type": row[3],
+            "chunk_text": row[4],
+            "chunk_index": row[5]
+        }
+        for row in cursor.fetchall()
+    ]
+    conn.close()
+    return chunks
+
+def get_all_emails() -> list:
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT message_id, subject, sender, snippet, received_at FROM emails")
+    emails = [
+        {
+            "message_id": row[0],
+            "subject": row[1],
+            "sender": row[2],
+            "snippet": row[3],
+            "received_at": row[4]
+        }
+        for row in cursor.fetchall()
+    ]
+    conn.close()
+    return emails
+
+def insert_fallback_node(node_id: str, label: str, name: str):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        "INSERT OR REPLACE INTO graph_nodes (id, label, name) VALUES (?, ?, ?)",
+        (node_id, label, name)
+    )
+    conn.commit()
+    conn.close()
+
+def insert_fallback_relationship(rel_id: str, source_id: str, target_id: str, type_str: str):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        "INSERT OR REPLACE INTO graph_relationships (id, source_id, target_id, type) VALUES (?, ?, ?, ?)",
+        (rel_id, source_id, target_id, type_str)
+    )
+    conn.commit()
+    conn.close()
+
+def get_fallback_relationships(node_id: str) -> list:
+    conn = get_connection()
+    cursor = conn.cursor()
+    # Find relationships where node_id is either source or target
+    cursor.execute(
+        """
+        SELECT r.type, s.label, s.name, t.label, t.name
+        FROM graph_relationships r
+        JOIN graph_nodes s ON r.source_id = s.id
+        JOIN graph_nodes t ON r.target_id = t.id
+        WHERE r.source_id = ? OR r.target_id = ?
+        """,
+        (node_id, node_id)
+    )
+    results = [
+        {
+            "type": row[0],
+            "source_label": row[1],
+            "source_name": row[2],
+            "target_label": row[3],
+            "target_name": row[4]
+        }
+        for row in cursor.fetchall()
+    ]
+    conn.close()
+    return results
+
+def clear_fallback_graph():
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("DELETE FROM graph_relationships")
+    cursor.execute("DELETE FROM graph_nodes")
+    conn.commit()
+    conn.close()
+
+def delete_data_before_date(date_str: str) -> dict:
+    conn = get_connection()
+    cursor = conn.cursor()
+    
+    # We compare string-wise since ISO dates sort correctly
+    cursor.execute("DELETE FROM emails WHERE received_at < ?", (date_str,))
+    emails_deleted = cursor.rowcount
+    
+    cursor.execute("DELETE FROM repository_documents WHERE synced_at < ?", (date_str,))
+    docs_deleted = cursor.rowcount
+    
+    cursor.execute("DELETE FROM document_chunks WHERE created_at < ?", (date_str,))
+    chunks_deleted = cursor.rowcount
+    
+    conn.commit()
+    conn.close()
+    
+    return {
+        "emails": emails_deleted,
+        "documents": docs_deleted,
+        "chunks": chunks_deleted
+    }