core/database.py

import chromadb
from chromadb.config import Settings
from core.config import settings
import os

class Database:
    def __init__(self):
        self.client = chromadb.PersistentClient(
            path=settings.CHROMA_DB_PATH,
            settings=Settings(allow_reset=True)
        )
        
    def get_collection(self, name: str):
        """Get or create a collection"""
        try:
            return self.client.get_collection(name)
        except:
            return self.client.create_collection(name)
    
    def store_artifact(self, collection_name: str, document: str, metadata: dict):
        """Store a project artifact in the database"""
        collection = self.get_collection(collection_name)
        collection.add(
            documents=[document],
            metadatas=[metadata],
            ids=[f"doc_{len(collection.get()['ids'])}"]
        )
    
    def query_artifacts(self, collection_name: str, query: str, n_results: int = 5):
        """Query artifacts in a collection"""
        collection = self.get_collection(collection_name)
        results = collection.query(
            query_texts=[query],
            n_results=n_results
        )
        return results

db = Database()