Add API server setup, authentication, and RAG pipeline integration

- Updated README.md to include instructions for starting the API server and setting up authentication.
- Introduced debug_import.py for testing import paths and module availability.
- Added SQLModel and aiosqlite dependencies to pyproject.toml.
- Implemented database initialization and RAG pipeline setup in app.py.
- Created authentication logic in auth.py, including API key verification.
- Configured API settings in config.py to include an optional API key.
- Developed database connection and session management in database.py.
- Defined database models for analysis and findings in db_models.py.
- Established API routes for analysis and RAG functionalities in routes/analysis.py and routes/rag.py.
- Implemented in-memory and PostgreSQL vector store interfaces for RAG in store/base.py, store/memory.py, and store/postgres.py.
- Added tests for API functionality and authentication in test_api.py and test_auth.py.
- Updated pytest.ini to include necessary python paths for testing.
- Modified requirements.txt to include SQLModel and aiosqlite dependencies.

Author: moshesham

README.md

@@ -163,6 +163,35 @@ python -m omniscient_llm pull codellama:7b-instruct
 python -m omniscient_llm recommend --category code
 ```
 
+### API Server
+
+Start the REST API server:
+
+```bash
+# Ensure packages are in PYTHONPATH
+export PYTHONPATH=$PYTHONPATH:packages/api/src:packages/core/src:packages/rag/src:packages/llm/src
+
+# Run the server
+python -m omniscient_api.cli serve
+```
+
+The API will be available at `http://localhost:8000`.
+Documentation is available at `http://localhost:8000/docs`.
+
+#### Authentication
+
+The API is protected by an API Key. Set the `OMNISCIENT_API_KEY` environment variable to enable authentication.
+
+```bash
+export OMNISCIENT_API_KEY="your-secret-key"
+```
+
+Pass the key in the `X-API-Key` header:
+
+```bash
+curl -H "X-API-Key: your-secret-key" http://localhost:8000/api/v1/analyze ...
+```
+
 ---
 
 ## 🐳 Docker Deployment

debug_import.py

@@ -0,0 +1,18 @@
+import sys
+import os
+
+# Add paths
+sys.path.append(os.path.abspath("packages/api/src"))
+sys.path.append(os.path.abspath("packages/rag/src"))
+
+try:
+    from omniscient_api.app import create_app
+    print("Import create_app successful")
+    from omniscient_api.auth import verify_api_key
+    print("Import verify_api_key successful")
+    import tests.test_auth
+    print("Import tests.test_auth successful")
+except Exception as e:
+    print(f"Import failed: {e}")
+    import traceback
+    traceback.print_exc()

packages/api/pyproject.toml

@@ -34,6 +34,8 @@ dependencies = [
     "uvicorn>=0.24.0",
     "pydantic>=2.0.0",
     "python-multipart>=0.0.6",
+    "sqlmodel>=0.0.14",
+    "aiosqlite>=0.19.0",
 ]
 
 [project.optional-dependencies]

packages/api/src/omniscient_api/app.py

@@ -10,7 +10,9 @@
 from omniscient_core.logging import get_logger
 from .config import APIConfig, load_api_config
 from .routes import router as api_router
+from .routes.rag import router as rag_router
 from .models import HealthResponse
+from .database import init_db
 
 logger = get_logger(__name__)
 
@@ -26,12 +28,46 @@ async def lifespan(app: FastAPI):
 
     logger.info("Starting Omniscient API server")
 
-    # Initialize services
-    # TODO: Initialize LLM client, cache, etc.
+    # Initialize database
+    await init_db()
+    
+    # Initialize RAG Pipeline
+    try:
+        from omniscient_rag.pipeline import RAGPipeline
+        from omniscient_llm.providers.ollama import OllamaProvider
+        
+        config: APIConfig = app.state.config
+        
+        # Initialize LLM provider
+        llm_provider = OllamaProvider(
+            base_url=config.get_llm_base_url(),
+            model=config.llm.model,
+        )
+        
+        # Initialize RAG
+        # Use memory:// if no DATABASE_URL provided or if explicitly requested
+        import os
+        db_url = os.getenv("DATABASE_URL", "memory://")
+        
+        pipeline = RAGPipeline.create(
+            db_url=db_url,
+            embed_fn=llm_provider,
+            llm_fn=llm_provider.generate,
+        )
+        await pipeline.initialize()
+        app.state.rag_pipeline = pipeline
+        logger.info(f"RAG Pipeline initialized with {db_url}")
+        
+    except Exception as e:
+        logger.error(f"Failed to initialize RAG pipeline: {e}")
+        app.state.rag_pipeline = None
 
     yield
 
     # Cleanup
+    if app.state.rag_pipeline:
+        await app.state.rag_pipeline.close()
+    
     logger.info("Shutting down Omniscient API server")
 
 
@@ -97,6 +133,7 @@ async def root():
 
     # Include API routes
     app.include_router(api_router, prefix=config.api_prefix)
+    app.include_router(rag_router, prefix=config.api_prefix)
 
     logger.info(f"API configured at {config.host}:{config.port}{config.api_prefix}")
 

packages/api/src/omniscient_api/auth.py

@@ -0,0 +1,32 @@
+"""Authentication dependencies."""
+
+from typing import Optional
+from fastapi import Security, HTTPException, status
+from fastapi.security import APIKeyHeader
+from .config import get_config
+
+api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+async def verify_api_key(
+    api_key: str = Security(api_key_header),
+) -> Optional[str]:
+    """Verify API key if configured."""
+    config = get_config()
+    
+    # If no API key is configured, allow all requests
+    if not config.api_key:
+        return api_key
+
+    if not api_key:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Missing API Key",
+        )
+
+    if api_key != config.api_key:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Invalid API Key",
+        )
+        
+    return api_key

packages/api/src/omniscient_api/config.py

@@ -55,6 +55,7 @@ class APIConfig(BaseSettings):
     port: int = Field(default=8000, alias="OMNISCIENT_API_PORT")
     workers: int = Field(default=1, alias="OMNISCIENT_API_WORKERS")
     debug: bool = Field(default=False, alias="OMNISCIENT_DEBUG")
+    api_key: Optional[str] = Field(default=None, description="API Key for authentication")
 
     # API settings
     api_prefix: str = "/api/v1"
@@ -76,8 +77,8 @@ class APIConfig(BaseSettings):
 
     model_config = SettingsConfigDict(
         env_prefix="OMNISCIENT_",
-        env_file=".env",
-        extra="ignore",
+        env_nested_delimiter="__",
+        case_sensitive=False
     )
 
     def get_llm_base_url(self) -> str:
@@ -132,3 +133,11 @@ def load_api_config(config_path: Optional[str] = None) -> APIConfig:
                     config.llm.model = llm["model"]
 
     return config
+
+
+from functools import lru_cache
+
+@lru_cache()
+def get_config() -> APIConfig:
+    """Get cached API configuration."""
+    return load_api_config()

packages/api/src/omniscient_api/database.py

@@ -0,0 +1,20 @@
+from typing import AsyncGenerator
+from sqlmodel import SQLModel
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+
+DATABASE_URL = "sqlite+aiosqlite:///./omniscient.db"
+
+engine = create_async_engine(DATABASE_URL, echo=True, future=True)
+
+async def init_db():
+    async with engine.begin() as conn:
+        # await conn.run_sync(SQLModel.metadata.drop_all)
+        await conn.run_sync(SQLModel.metadata.create_all)
+
+async def get_session() -> AsyncGenerator[AsyncSession, None]:
+    async_session = sessionmaker(
+        engine, class_=AsyncSession, expire_on_commit=False
+    )
+    async with async_session() as session:
+        yield session

packages/api/src/omniscient_api/db_models.py

@@ -0,0 +1,39 @@
+from typing import Optional, List
+from datetime import datetime
+from sqlmodel import SQLModel, Field, Relationship, JSON
+from omniscient_core import AnalysisStatus, Severity, FindingCategory
+import uuid
+
+class AnalysisBase(SQLModel):
+    repository_url: str
+    branch: Optional[str] = None
+    status: AnalysisStatus = AnalysisStatus.PENDING
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    error: Optional[str] = None
+    # Storing summary and metrics as JSON for simplicity for now
+    summary: Optional[dict] = Field(default=None, sa_type=JSON)
+    metrics: Optional[dict] = Field(default=None, sa_type=JSON)
+
+class Analysis(AnalysisBase, table=True):
+    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()), primary_key=True)
+    findings: List["FindingModel"] = Relationship(back_populates="analysis")
+
+class FindingModel(SQLModel, table=True):
+    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()), primary_key=True)
+    analysis_id: str = Field(foreign_key="analysis.id")
+    
+    agent_name: str
+    severity: Severity
+    category: FindingCategory
+    title: str
+    description: str
+    file_path: Optional[str] = None
+    line_start: Optional[int] = None
+    line_end: Optional[int] = None
+    code_snippet: Optional[str] = None
+    suggestions: List[str] = Field(default_factory=list, sa_type=JSON)
+    metadata_json: dict = Field(default_factory=dict, sa_type=JSON, alias="metadata")
+
+    analysis: Analysis = Relationship(back_populates="findings")

packages/api/src/omniscient_api/routes/__init__.py

@@ -0,0 +1,3 @@
+"""Routes package."""
+
+from .analysis import router