53295444de
Features: - Fixed Ollama configuration for native (localhost) instead of Docker - Increased LLM timeout from 30s to 120s for quality answers - Implemented KB slug system for clean MCP endpoint URLs - Added multiple knowledge base support with type metadata - Integrated MCP server with FastAPI (port 8051) - React frontend with KB management UI - PostgreSQL + pgvector for vector embeddings Performance: - 33-56% faster than Docker setup with native Ollama - Direct Metal GPU acceleration on Apple Silicon - Optimized RAG engine with caching - ~40s response time for comprehensive answers Documentation: - Complete native setup guide (README_NATIVE_SETUP.md) - Performance benchmarking and comparison - MCP endpoint documentation - Slug implementation details - Decision rationale for native vs Docker Cleanup: - Removed old Streamlit setup files - Consolidated setup scripts to single setup_native.sh - Consolidated service management scripts - Removed test artifacts and old analysis docs - Ready for sharing with team 🚀 Production-ready for local team deployment 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com>
542 lines
20 KiB
Python
542 lines
20 KiB
Python
"""
|
|
PostgreSQL + pgvector database manager for Code-Forge.
|
|
Replaces ChromaDB with a more robust, ACID-compliant solution.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import re
|
|
from typing import List, Dict, Any, Optional, Tuple
|
|
from datetime import datetime
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor, Json
|
|
from psycopg2.pool import SimpleConnectionPool
|
|
import numpy as np
|
|
|
|
from app.core.config import Config
|
|
from app.core.kb_types import KBType, KBMetadata
|
|
|
|
|
|
def generate_slug(name: str) -> str:
|
|
"""
|
|
Generate a URL-friendly slug from a KB name.
|
|
|
|
Examples:
|
|
"Pistn Agent OS" -> "pistn-agent-os"
|
|
"My Code Base!" -> "my-code-base"
|
|
"Test_KB 123" -> "test-kb-123"
|
|
"""
|
|
# Convert to lowercase
|
|
slug = name.lower()
|
|
# Replace spaces and underscores with hyphens
|
|
slug = re.sub(r'[\s_]+', '-', slug)
|
|
# Remove any characters that aren't alphanumeric or hyphens
|
|
slug = re.sub(r'[^a-z0-9-]', '', slug)
|
|
# Remove multiple consecutive hyphens
|
|
slug = re.sub(r'-+', '-', slug)
|
|
# Remove leading/trailing hyphens
|
|
slug = slug.strip('-')
|
|
return slug
|
|
|
|
|
|
class PostgresManager:
|
|
"""Manages PostgreSQL database operations with pgvector for embeddings."""
|
|
|
|
def __init__(self):
|
|
"""Initialize connection pool."""
|
|
# Build connection parameters
|
|
conn_params = {
|
|
"minconn": 1,
|
|
"maxconn": 10,
|
|
"host": os.getenv("POSTGRES_HOST", "localhost"),
|
|
"port": int(os.getenv("POSTGRES_PORT", "5432")),
|
|
"database": os.getenv("POSTGRES_DB", "codeforge"),
|
|
"user": os.getenv("POSTGRES_USER", os.getenv("USER", "postgres")),
|
|
}
|
|
|
|
# Only add password if it's set (for local PostgreSQL with trust auth)
|
|
password = os.getenv("POSTGRES_PASSWORD", "")
|
|
if password:
|
|
conn_params["password"] = password
|
|
|
|
self.pool = SimpleConnectionPool(**conn_params)
|
|
self._init_schema()
|
|
|
|
def _init_schema(self):
|
|
"""Initialize database schema from schema.sql."""
|
|
# Schema is managed externally - just verify connection
|
|
# The schema.sql should be run manually on the local PostgreSQL
|
|
pass
|
|
|
|
def get_connection(self):
|
|
"""Get a connection from the pool."""
|
|
return self.pool.getconn()
|
|
|
|
def return_connection(self, conn):
|
|
"""Return a connection to the pool."""
|
|
self.pool.putconn(conn)
|
|
|
|
# Knowledge Base Operations
|
|
|
|
def create_collection(
|
|
self,
|
|
name: str,
|
|
kb_type: KBType = KBType.GENERIC,
|
|
description: str = "",
|
|
tags: List[str] = None,
|
|
embed_dim: int = 768
|
|
) -> Dict[str, Any]:
|
|
"""Create a new knowledge base with PGVectorStore table."""
|
|
import json
|
|
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
# Generate URL-friendly slug
|
|
slug = generate_slug(name)
|
|
|
|
# Sanitize name for table name
|
|
table_name = f"data_{name.lower().replace(' ', '_').replace('-', '_')}"
|
|
|
|
# Prepare metadata (tags go in metadata JSON, not separate column)
|
|
metadata = {}
|
|
if tags:
|
|
metadata["tags"] = tags
|
|
|
|
# Insert KB metadata (no tags column in NEW schema)
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO knowledge_bases (name, slug, kb_type, description, metadata, table_name, embed_dim)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
|
RETURNING id, name, slug, kb_type, description, created_at, metadata, table_name, embed_dim
|
|
""",
|
|
(name, slug, kb_type.value, description, json.dumps(metadata), table_name, embed_dim)
|
|
)
|
|
result = cur.fetchone()
|
|
|
|
# Create the PGVectorStore table using the helper function
|
|
cur.execute("SELECT create_kb_table(%s, %s, %s)", (name, table_name, embed_dim))
|
|
|
|
conn.commit()
|
|
return dict(result)
|
|
except psycopg2.IntegrityError as e:
|
|
conn.rollback()
|
|
if 'slug' in str(e):
|
|
raise ValueError(f"A knowledge base with a similar name already exists (slug conflict: '{slug}')")
|
|
raise ValueError(f"Knowledge base '{name}' already exists")
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def delete_collection(self, name: str) -> bool:
|
|
"""Delete a knowledge base and its PGVectorStore table."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
# Get table name before deleting
|
|
cur.execute("SELECT table_name FROM knowledge_bases WHERE name = %s", (name,))
|
|
result = cur.fetchone()
|
|
|
|
if not result:
|
|
return False
|
|
|
|
table_name = result['table_name']
|
|
|
|
# Drop the PGVectorStore table if it exists
|
|
if table_name:
|
|
cur.execute("SELECT drop_kb_table(%s)", (table_name,))
|
|
|
|
# Delete KB metadata
|
|
cur.execute("DELETE FROM knowledge_bases WHERE name = %s", (name,))
|
|
deleted = cur.rowcount > 0
|
|
|
|
conn.commit()
|
|
return deleted
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def list_collections(self) -> List[Dict[str, Any]]:
|
|
"""List all knowledge bases with metadata."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT name, slug, kb_type, description, created_at, metadata
|
|
FROM knowledge_bases
|
|
ORDER BY created_at DESC
|
|
"""
|
|
)
|
|
results = cur.fetchall()
|
|
|
|
return [
|
|
{
|
|
"name": row["name"],
|
|
"slug": row.get("slug") or generate_slug(row["name"]), # Fallback for old KBs
|
|
"metadata": {
|
|
"kb_type": row["kb_type"],
|
|
"description": row["description"] or "",
|
|
"created_at": row["created_at"].isoformat() if row["created_at"] else "",
|
|
**row["metadata"]
|
|
}
|
|
}
|
|
for row in results
|
|
]
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def get_collection_metadata(self, name: str) -> Dict[str, Any]:
|
|
"""Get metadata for a specific knowledge base."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT kb_type, description, created_at, tags, metadata
|
|
FROM knowledge_bases
|
|
WHERE name = %s
|
|
""",
|
|
(name,)
|
|
)
|
|
row = cur.fetchone()
|
|
|
|
if not row:
|
|
raise ValueError(f"Knowledge base '{name}' not found")
|
|
|
|
return {
|
|
"kb_type": row["kb_type"],
|
|
"description": row["description"] or "",
|
|
"created_at": row["created_at"].isoformat() if row["created_at"] else "",
|
|
"tags": row["tags"] or "",
|
|
**row["metadata"]
|
|
}
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def list_collections_by_type(self, kb_type: KBType) -> List[Dict[str, Any]]:
|
|
"""List knowledge bases filtered by type."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT name, kb_type, description, created_at, metadata
|
|
FROM knowledge_bases
|
|
WHERE kb_type = %s
|
|
ORDER BY created_at DESC
|
|
""",
|
|
(kb_type.value,)
|
|
)
|
|
results = cur.fetchall()
|
|
|
|
return [
|
|
{
|
|
"name": row["name"],
|
|
"metadata": {
|
|
"kb_type": row["kb_type"],
|
|
"description": row["description"] or "",
|
|
"created_at": row["created_at"].isoformat() if row["created_at"] else "",
|
|
**row["metadata"]
|
|
}
|
|
}
|
|
for row in results
|
|
]
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
# Document Operations
|
|
|
|
def add_documents(
|
|
self,
|
|
kb_name: str,
|
|
documents: List[str],
|
|
embeddings: List[List[float]],
|
|
metadatas: List[Dict[str, Any]],
|
|
ids: List[str]
|
|
) -> None:
|
|
"""Add documents to a knowledge base using PGVectorStore schema."""
|
|
from llama_index.core.schema import TextNode
|
|
from llama_index.vector_stores.postgres import PGVectorStore
|
|
|
|
conn = self.get_connection()
|
|
try:
|
|
# Get KB metadata
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(
|
|
"SELECT id, table_name, embed_dim FROM knowledge_bases WHERE name = %s",
|
|
(kb_name,)
|
|
)
|
|
result = cur.fetchone()
|
|
if not result:
|
|
raise ValueError(f"Knowledge base '{kb_name}' not found")
|
|
|
|
kb_id = result['id']
|
|
table_name = result['table_name']
|
|
embed_dim = result['embed_dim']
|
|
|
|
if not table_name:
|
|
raise ValueError(f"Knowledge base '{kb_name}' has not been migrated to PGVectorStore schema")
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
# Create PGVectorStore instance
|
|
pg_user = os.getenv("POSTGRES_USER", os.getenv("USER", "postgres"))
|
|
pg_password = os.getenv("POSTGRES_PASSWORD", "")
|
|
pg_host = os.getenv("POSTGRES_HOST", "localhost")
|
|
pg_port = os.getenv("POSTGRES_PORT", "5432")
|
|
pg_db = os.getenv("POSTGRES_DB", "codeforge")
|
|
|
|
# PGVectorStore automatically adds 'data_' prefix to table names
|
|
# Our schema already has 'data_' prefix, so we need to strip it
|
|
pgvector_table_name = table_name.removeprefix('data_')
|
|
|
|
# Use individual parameters to avoid async connection string issues
|
|
vector_store = PGVectorStore.from_params(
|
|
host=pg_host,
|
|
port=pg_port,
|
|
database=pg_db,
|
|
user=pg_user,
|
|
password=pg_password,
|
|
table_name=pgvector_table_name, # PGVectorStore will add 'data_' prefix
|
|
embed_dim=embed_dim,
|
|
perform_setup=False # Don't recreate tables
|
|
)
|
|
|
|
# Create TextNode objects for llama_index
|
|
nodes = []
|
|
for doc, emb, meta, doc_id in zip(documents, embeddings, metadatas, ids):
|
|
# Add kb_id to metadata
|
|
meta_with_kb = {**meta, "kb_id": str(kb_id)}
|
|
|
|
node = TextNode(
|
|
text=doc,
|
|
id_=doc_id,
|
|
embedding=emb,
|
|
metadata=meta_with_kb
|
|
)
|
|
nodes.append(node)
|
|
|
|
# Add nodes to vector store
|
|
vector_store.add(nodes)
|
|
|
|
def get_documents_by_metadata(
|
|
self,
|
|
kb_name: str,
|
|
where: Dict[str, Any],
|
|
limit: Optional[int] = None
|
|
) -> List[Dict[str, Any]]:
|
|
"""Get documents filtered by metadata from NEW PGVectorStore schema."""
|
|
conn = self.get_connection()
|
|
try:
|
|
# Get table name for this KB
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT table_name FROM knowledge_bases WHERE name = %s", (kb_name,))
|
|
result = cur.fetchone()
|
|
if not result or not result[0]:
|
|
return []
|
|
table_name = result[0]
|
|
|
|
# Build WHERE clause for JSONB metadata_ column
|
|
where_conditions = []
|
|
params = []
|
|
for key, value in where.items():
|
|
where_conditions.append(f"metadata_->>%s = %s")
|
|
params.extend([key, str(value)])
|
|
|
|
where_clause = " AND ".join(where_conditions) if where_conditions else "TRUE"
|
|
|
|
# Query NEW schema: data_{kb_name} table with columns: id, text, metadata_, node_id, embedding
|
|
query = f"""
|
|
SELECT node_id, text, metadata_
|
|
FROM {table_name}
|
|
WHERE {where_clause}
|
|
"""
|
|
if limit:
|
|
query += f" LIMIT {limit}"
|
|
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(query, params)
|
|
results = cur.fetchall()
|
|
# Return with 'metadata' key (not 'metadata_') for compatibility
|
|
return [
|
|
{
|
|
"node_id": row["node_id"],
|
|
"text": row["text"],
|
|
"metadata": row["metadata_"]
|
|
}
|
|
for row in results
|
|
]
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def query_documents(
|
|
self,
|
|
kb_name: str,
|
|
query_embedding: List[float],
|
|
n_results: int = 10,
|
|
where: Optional[Dict[str, Any]] = None
|
|
) -> Dict[str, List]:
|
|
"""Query documents using vector similarity search."""
|
|
conn = self.get_connection()
|
|
try:
|
|
# Get KB ID
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT id FROM knowledge_bases WHERE name = %s", (kb_name,))
|
|
result = cur.fetchone()
|
|
if not result:
|
|
raise ValueError(f"Knowledge base '{kb_name}' not found")
|
|
kb_id = result[0]
|
|
|
|
# Search documents
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT * FROM search_documents(%s, %s, %s, %s)
|
|
""",
|
|
(kb_id, query_embedding, n_results, Json(where) if where else None)
|
|
)
|
|
results = cur.fetchall()
|
|
|
|
return {
|
|
"ids": [[row["doc_id"] for row in results]],
|
|
"documents": [[row["content"] for row in results]],
|
|
"metadatas": [[row["metadata"] for row in results]],
|
|
"distances": [[1 - row["similarity"] for row in results]]
|
|
}
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def collection_exists(self, kb_name: str) -> bool:
|
|
"""Check if a knowledge base exists by name."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT EXISTS(SELECT 1 FROM knowledge_bases WHERE name = %s)", (kb_name,))
|
|
return cur.fetchone()[0]
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def get_kb_by_slug(self, slug: str) -> Optional[str]:
|
|
"""Get KB name by slug. Returns None if not found."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT name FROM knowledge_bases WHERE slug = %s", (slug,))
|
|
result = cur.fetchone()
|
|
return result[0] if result else None
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def slug_exists(self, slug: str) -> bool:
|
|
"""Check if a slug exists."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT EXISTS(SELECT 1 FROM knowledge_bases WHERE slug = %s)", (slug,))
|
|
return cur.fetchone()[0]
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def get_collection_count(self, kb_name: str) -> int:
|
|
"""Get document count for a knowledge base."""
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
# Get table name
|
|
cur.execute("SELECT table_name FROM knowledge_bases WHERE name = %s", (kb_name,))
|
|
result = cur.fetchone()
|
|
|
|
if not result or not result['table_name']:
|
|
return 0
|
|
|
|
table_name = result['table_name']
|
|
|
|
# Count documents in PGVectorStore table
|
|
cur.execute(f"SELECT COUNT(*) as count FROM {table_name}")
|
|
count_result = cur.fetchone()
|
|
return count_result['count'] if count_result else 0
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
def close(self):
|
|
"""Close all connections in the pool."""
|
|
if self.pool:
|
|
self.pool.closeall()
|
|
def query_similar(
|
|
self,
|
|
kb_id: int,
|
|
query_embedding: List[float],
|
|
top_k: int = 5
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Query for similar documents using vector similarity.
|
|
Works with NEW PGVectorStore schema (data_{kb_name} tables).
|
|
|
|
Args:
|
|
kb_id: Knowledge base ID
|
|
query_embedding: Query embedding vector
|
|
top_k: Number of results to return
|
|
|
|
Returns:
|
|
List of similar documents with text, metadata, and similarity score
|
|
"""
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
conn = self.get_connection()
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
# Get table name for this KB
|
|
cur.execute("SELECT table_name FROM knowledge_bases WHERE id = %s", (kb_id,))
|
|
result = cur.fetchone()
|
|
if not result or not result['table_name']:
|
|
logger.warning(f"No table found for kb_id={kb_id}")
|
|
return []
|
|
|
|
table_name = result['table_name']
|
|
|
|
# Convert embedding list to PostgreSQL vector format string
|
|
# Format: '[0.123, 0.456, ...]'
|
|
embedding_str = "[" + ",".join(map(str, query_embedding)) + "]"
|
|
logger.info(f"Query embedding dimension: {len(query_embedding)}")
|
|
logger.info(f"Querying kb_id={kb_id}, table={table_name}, top_k={top_k}")
|
|
|
|
# Query for similar documents using cosine similarity
|
|
# NEW schema uses 'text' and 'metadata_' columns
|
|
query = f"""
|
|
SELECT
|
|
text,
|
|
metadata_ as metadata,
|
|
node_id,
|
|
1 - (embedding <=> %s::vector) as similarity
|
|
FROM {table_name}
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT %s
|
|
"""
|
|
logger.info(f"Executing query with params: top_k={top_k}")
|
|
|
|
# Pass the embedding as a string with ::vector cast
|
|
cur.execute(query, (embedding_str, embedding_str, top_k))
|
|
|
|
results = cur.fetchall()
|
|
logger.info(f"PostgreSQL returned {len(results)} results")
|
|
if results:
|
|
logger.info(f"First result similarity: {results[0].get('similarity')}")
|
|
return [dict(row) for row in results]
|
|
except Exception as e:
|
|
logger.error(f"Error in query_similar: {e}")
|
|
raise
|
|
finally:
|
|
self.return_connection(conn)
|
|
|
|
|
|
# Singleton instance
|
|
_pg_manager = None
|
|
|
|
|
|
def get_pg_manager() -> PostgresManager:
|
|
"""Get or create the PostgresManager singleton."""
|
|
global _pg_manager
|
|
if _pg_manager is None:
|
|
_pg_manager = PostgresManager()
|
|
return _pg_manager
|
|
|