Initial commit: add .gitignore and README

fusionagi/memory/vector_pgvector.py

@@ -0,0 +1,101 @@
+"""pgvector-backed VectorMemory adapter. Requires: pip install fusionagi[vector]."""
+
+from typing import Any
+
+from fusionagi._logger import logger
+
+
+def create_vector_memory_pgvector(
+    connection_string: str,
+    table_name: str = "embeddings",
+    dimension: int = 1536,
+) -> Any:
+    """
+    Create pgvector-backed VectorMemory when pgvector is installed.
+    Returns None if pgvector/database unavailable.
+    """
+    try:
+        import pgvector
+        from pgvector.psycopg import register_vector
+    except ImportError:
+        logger.debug("pgvector not installed; use pip install fusionagi[vector]")
+        return None
+
+    try:
+        import psycopg
+    except ImportError:
+        logger.debug("psycopg not installed; use pip install fusionagi[memory]")
+        return None
+
+    return VectorMemoryPgvector(connection_string, table_name, dimension)
+
+
+class VectorMemoryPgvector:
+    """VectorMemory implementation using pgvector."""
+
+    def __init__(
+        self,
+        connection_string: str,
+        table_name: str = "embeddings",
+        dimension: int = 1536,
+    ) -> None:
+        import pgvector
+        from pgvector.psycopg import register_vector
+
+        self._conn_str = connection_string
+        self._table = table_name
+        self._dim = dimension
+
+        with psycopg.connect(connection_string) as conn:
+            register_vector(conn)
+            with conn.cursor() as cur:
+                cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+                cur.execute(
+                    f"""
+                    CREATE TABLE IF NOT EXISTS {table_name} (
+                        id TEXT PRIMARY KEY,
+                        embedding vector({dimension}),
+                        metadata JSONB,
+                        created_at TIMESTAMPTZ DEFAULT NOW()
+                    )
+                    """
+                )
+            conn.commit()
+
+    def add(self, id: str, embedding: list[float], metadata: dict[str, Any] | None = None) -> None:
+        import json
+        import psycopg
+        from pgvector.psycopg import register_vector
+
+        with psycopg.connect(self._conn_str) as conn:
+            register_vector(conn)
+            with conn.cursor() as cur:
+                cur.execute(
+                    f"""
+                    INSERT INTO {self._table} (id, embedding, metadata)
+                    VALUES (%s, %s, %s)
+                    ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding, metadata = EXCLUDED.metadata
+                    """,
+                    (id, embedding, json.dumps(metadata or {})),
+                )
+            conn.commit()
+
+    def search(self, query_embedding: list[float], top_k: int = 10) -> list[dict[str, Any]]:
+        import json
+        import psycopg
+        from pgvector.psycopg import register_vector
+
+        with psycopg.connect(self._conn_str) as conn:
+            register_vector(conn)
+            with conn.cursor() as cur:
+                cur.execute(
+                    f"""
+                    SELECT id, metadata
+                    FROM {self._table}
+                    ORDER BY embedding <-> %s
+                    LIMIT %s
+                    """,
+                    (query_embedding, top_k),
+                )
+                rows = cur.fetchall()
+        return [{"id": r[0], "metadata": json.loads(r[1]) if r[1] else {}} for r in rows]