[feat] add chroma conector

run.py

@@ -4,4 +4,4 @@
 app = create_app()
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8001)

src/api/controllers/__init__.py

@@ -1,4 +1,5 @@
 from .chat import new_message as chat_new_message
 from .guardrails import Guardrail
+from .controller_document import add_documents as add_documents
 
-__all__ = ["chat_new_message", "Guardrail"]
+__all__ = ["chat_new_message", "add_documents", "Guardrail"]

src/api/controllers/controller_document.py

@@ -0,0 +1,61 @@
+from src.infrastructure.database.chromadb.conector import ChromaDB
+from src.services.docmuent_extration.extractor import DocumentFileReader
+from typing import List, Dict, Optional
+
+
+async def add_documents(
+    db: ChromaDB,
+    file_path: str,
+    collection_name: str,
+    metadata: Optional[str] = None,
+    documnet_id: Optional[List[str]] = None
+) -> Dict:
+
+    reader = DocumentFileReader(file_path, metadata, documnet_id)
+    result = await db.add_documents(
+        documents=reader.documents,
+        collection_name=collection_name,
+        metadatas=reader.metadata,
+        ids=reader.ids
+    )
+    return result
+
+async def query_documents(
+    db: ChromaDB,
+    query_text: str,
+    collection_name: str,
+    n_results: int = 5,
+    where: Optional[dict] = None
+):
+    result = await db.query_documents(
+        query_text=query_text,
+        collection_name=collection_name,
+        n_results=n_results,
+        where=where) 
+    return result
+
+async def delete_documents(
+    db: ChromaDB,
+    ids: List[str],
+    collection_name: str
+):
+    result = await db.delete_documents(
+        ids=ids,
+        collection_name=collection_name
+    )
+    return result
+
+async def list_collections(
+    db: ChromaDB
+):
+    result = await db.list_collections()
+    return result
+
+async def list_documents(
+    db: ChromaDB,
+    collection_name: str,
+):
+    result = await db.list_documents(
+        collection_name=collection_name
+    )
+    return result

src/api/models/__init__.py

@@ -1,3 +1,4 @@
 from .api import APIResponse, APIRequest
+from .documents_request import AddDocumentRequest, QueryDocumentRequest, DeleteDocumentRequest
 
-__all__ = ["APIResponse", "APIRequest"]
+__all__ = ["APIResponse", "APIRequest", "AddDocumentRequest", "QueryDocumentRequest", "DeleteDocumentRequest"]

src/api/models/documents_request.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel
+from typing import Optional, List, Dict
+
+class AddDocumentRequest(BaseModel):
+    file_path: str
+    collection_name: str
+
+class QueryDocumentRequest(BaseModel):
+    query_text: str
+    collection_name: str
+    n_results: Optional[int] = 5
+    where: Optional[Dict] = None
+
+class DeleteDocumentRequest(BaseModel):
+    ids: List[str]
+    collection_name: str
+
+class ListDocumentsRequest(BaseModel):   
+    collection_name: str
+
+class ListCollectionsRequest(BaseModel):
+    pass

src/api/routes/__init__.py

@@ -0,0 +1,6 @@
+from src.api.routes.route_document import router as document_router
+from src.api.routes.chat import router as chat_router
+
+
+__all__ = ["chat_router", "document_router"]
+

src/api/routes/route_document.py

@@ -0,0 +1,99 @@
+from fastapi import APIRouter, status, Request, Depends, HTTPException
+from src.api.models import APIResponse
+from src.api.models.documents_request import AddDocumentRequest, QueryDocumentRequest, ListDocumentsRequest, ListCollectionsRequest, DeleteDocumentRequest
+from typing import Optional, Dict
+from src.api.controllers.controller_document import add_documents, query_documents, delete_documents, list_collections, list_documents
+
+router = APIRouter(
+    prefix="/documents",
+    tags=["documents"]
+)
+
+@router.post("/add_documents", status_code=status.HTTP_200_OK)
+async def add_document_route(doc_request: AddDocumentRequest, req: Request) -> APIResponse:
+    try:
+        result = await add_documents(
+            req.app.vectordb,
+            doc_request.file_path,
+            doc_request.collection_name
+        )
+
+        return APIResponse(
+            status_code=200,
+            response=result
+        )
+
+    except Exception as e:
+        return APIResponse(
+            status_code=500,
+            status_message=f"Error detalhado: {str(e)}"
+        )
+
+@router.post("/query_documents", status_code=status.HTTP_200_OK)
+async def query_document_route(doc_request: QueryDocumentRequest, req: Request) -> APIResponse:
+    try:
+        result = await query_documents(
+            req.app.vectordb,
+            doc_request.query_text,
+            doc_request.collection_name,
+            doc_request.n_results,
+            doc_request.where
+        )
+        return APIResponse(
+            status_code=200,
+            response=result
+        )
+    except Exception as e:
+        return APIResponse(
+            status_code=500,
+            status_message=f"Error detalhado: {str(e)}"
+        )
+
+@router.get("/list_collections", status_code=status.HTTP_200_OK)
+async def list_collections_route(req: Request):
+    try:
+        result = await list_collections(req.app.vectordb)
+        return {
+            "status_code": 200,
+            "response": result  # Retorna diretamente a lista de nomes das coleções
+        }
+    except Exception as e:
+        return {
+            "status_code": 500,
+            "status_message": f"Erro detalhado: {str(e)}"
+        }
+
+@router.post("/list_documents", status_code=status.HTTP_200_OK)
+async def list_documents_route(doc_request: ListDocumentsRequest, req: Request) -> APIResponse:
+    try:
+        result = await list_documents(
+            req.app.vectordb,
+            doc_request.collection_name
+        )
+        return APIResponse(
+            status_code=200,
+            response=result
+        )
+    except Exception as e:
+        return APIResponse(
+            status_code=500,
+            status_message=f"Error detalhado: {str(e)}"
+        )
+
+@router.delete("/delete_documents", status_code=status.HTTP_200_OK)
+async def delete_document_route(doc_request: DeleteDocumentRequest, req: Request) -> APIResponse:
+    try:
+        result = await delete_documents(
+            req.app.vectordb,
+            doc_request.ids,
+            doc_request.collection_name
+        )
+        return APIResponse(
+            status_code=200,
+            response=result
+        )
+    except Exception as e:
+        return APIResponse(
+            status_code=500,
+            status_message=f"Error detalhado: {str(e)}"
+        )

src/infrastructure/database/chromadb/conector.py

@@ -0,0 +1,125 @@
+import chromadb
+from chromadb.config import Settings
+from typing import List, Optional
+
+
+class ChromaDB:
+    """Manager for ChromaDB connection and operations."""
+
+    def __init__(self):
+        """Initialize ChromaDB connection."""
+        self.host = 'localhost'
+        self.port = 8000
+        self.client = self._connect()
+        self.collection = None
+
+    def _connect(self):
+        """Connect to ChromaDB."""
+        client = chromadb.HttpClient(host=self.host, port=self.port)
+        return client
+
+    def _get_or_create_collection(self, collection_name: str):
+        """Create or get an existing collection.
+
+        Args:
+            collection_name (str): Name of the collection to create/get
+
+        Returns:
+            Collection: Created/retrieved collection object
+        """
+        self.collection = self.client.get_or_create_collection(
+            name=collection_name
+        )
+        return self.collection
+
+    async def add_documents(
+        self,
+        documents: List[str],
+        collection_name: str,
+        metadatas: Optional[List[dict]] = None,
+        ids: Optional[List[str]] = None
+    ):
+        """Add documents to the collection.
+
+        Args:
+            documents (List[str]): List of document texts
+            metadatas (Optional[List[dict]]): Document metadata
+            ids (Optional[List[str]]): Unique document IDs
+
+        Raises:
+            ValueError: If collection is not initialized
+
+        Returns:
+            dict: Result of add operation
+        """
+
+        self.collection = self._get_or_create_collection(collection_name)
+        return self.collection.add(
+            documents=documents,
+            metadatas=metadatas,
+            ids=ids
+        )
+
+    async def query_documents(
+        self,
+        query_text: str,
+        collection_name: str,
+        n_results: int = 5,
+        where: Optional[dict] = None
+    ):
+        """Query similar documents in the collection.
+
+        Args:
+            query_text (str): Text to search for similarity
+            n_results (int): Number of desired results
+            where (Optional[dict]): Additional filters
+
+        Raises:
+            ValueError: If collection is not initialized
+
+        Returns:
+            dict: Query results
+        """
+
+        self.collection = self._get_or_create_collection(collection_name)
+        return self.collection.query(
+            query_texts=[query_text],
+            n_results=n_results,
+            where=where
+        )
+
+    async def list_collections(self):
+        """List all collections in ChromaDB.
+
+        Returns:
+            dict: List of collections
+        """
+        return self.client.list_collections()
+
+    async def list_documents(self, collection_name: str):
+        """List all documents in a
+        collection. Args:   collection_name (str): Name of the collection to list documents from Returns:     dict: List of documents in the collection"""
+
+        self.collection = self._get_or_create_collection(collection_name)
+        return self.collection.get(include=["documents", "metadatas"])
+
+    async def delete_documents(self, ids: List[str],collection_name: str):
+        """Delete documents from collection by IDs.
+
+        Args:
+            ids (List[str]): List of document IDs to delete
+
+        Raises:
+            ValueError: If collection is not initialized
+
+        Returns:
+            dict: Result of delete operation
+        """
+
+        self.collection = self._get_or_create_collection(collection_name)
+        return self.collection.delete(ids=ids)
+
+    def __del__(self):
+        """Close ChromaDB connection and clean up resources."""
+        if self.client:
+            self.client.reset()

src/main.py

@@ -1,21 +1,23 @@
 from fastapi import FastAPI
-
-from src.infrastructure.database import MongoDB
+from src.infrastructure.database.chromadb.conector import ChromaDB
+from src.infrastructure.database.mongodb.connector import MongoDB
 from src.infrastructure.config.llm import LLM
 from src.services.llama_guard import LlamaGuard
 
-from src.api.routes import chat_router
+from src.api.routes import chat_router, document_router
 
 
 def create_app():
     app = FastAPI()
 
     # defining API variables
+    app.vectordb = ChromaDB()
     app.database = MongoDB()
     app.llm = LLM()
     app.llama_guard = LlamaGuard()
 
     # including routes
     app.include_router(chat_router)
+    app.include_router(document_router)
 
     return app