diff --git a/bolna/ingestion_server/datachecks/__init__.py b/bolna/ingestion_server/datachecks/__init__.py index 7b0b27fe..3bede1da 100644 --- a/bolna/ingestion_server/datachecks/__init__.py +++ b/bolna/ingestion_server/datachecks/__init__.py @@ -46,4 +46,4 @@ class RAGTask(BaseModel): _status:str = RAGTaskStatus.WAIT _message:str = "" _index:str = "" - _nodes:list = [] + _nodes:list = [] \ No newline at end of file diff --git a/bolna/ingestion_server/main.py b/bolna/ingestion_server/main.py index f37ec5ba..e1e6c875 100644 --- a/bolna/ingestion_server/main.py +++ b/bolna/ingestion_server/main.py @@ -1,10 +1,13 @@ -from fastapi import FastAPI, UploadFile +import time +import json +import dotenv + +from fastapi import FastAPI, File, UploadFile, Form +from typing import Dict, Any, Optional, Union + from bolna.ingestion_server.datachecks import RAGConfig from bolna.ingestion_server import RAGFactory import uvicorn -import time -from typing import Dict -import dotenv dotenv.load_dotenv() DB: Dict[str, RAGConfig] = {} @@ -22,10 +25,8 @@ def heartbeat() -> float: @app.post("/create-rag") def create_rag(request: RAGConfig) -> Dict[str, str]: """Create a RAG configuration and return its ID. - Args: request (RAGConfig): The RAG configuration to create. - Returns: Dict[str, str]: A dictionary containing the created RAG ID. """ @@ -37,11 +38,9 @@ def create_rag(request: RAGConfig) -> Dict[str, str]: @app.post("/rag-upload-file/{rag_id}") async def rag_upload_file(file: UploadFile, rag_id: str) -> Dict[str, str]: """Upload a file for a specific RAG ID. - Args: file (UploadFile): The file to upload. rag_id (str): The ID of the RAG to associate with the file. - Returns: Dict[str, str]: A dictionary containing the upload status and index. """ @@ -50,21 +49,64 @@ async def rag_upload_file(file: UploadFile, rag_id: str) -> Dict[str, str]: return {"index": task._index, "status": task._status, "message": "DONE"} except Exception as e: return {"index": None, "status": "ERROR", "message": f"{e}"} + +@app.post("/make-rag") +async def make_rag( + file: UploadFile = File(...), + config: str = Form(...) +) -> Dict[str, Any]: + """ + Create a RAG configuration, return its ID, and ingest the uploaded file. + + Args: + file (UploadFile): The file to upload and ingest. + config (str): The RAG configuration as a JSON string. + Returns: + Dict[str, Any]: A dictionary containing the created RAG ID, upload status, and index. + """ + try: + # Parse the JSON string into a RAGConfig object + rag_config = RAGConfig.parse_raw(config) + + # Create RAG configuration + rag_id = rag_factory.make_rag(rag_config) + + # Ingest the file + task = await rag_factory.file_ingest(rag_name=rag_id, file=file) + + return { + "rag_id": rag_id, + "index": task._index, + "status": task._status, + "message": "RAG created and file ingested successfully" + } + except json.JSONDecodeError: + return { + "rag_id": None, + "index": None, + "status": "ERROR", + "message": "Invalid JSON in config parameter" + } + except Exception as e: + return { + "rag_id": None, + "index": None, + "status": "ERROR", + "message": f"Error creating RAG or ingesting file: {str(e)}" + } @app.get("/rag-retrive/{rag_id}/{index}") async def rag_retrive(query: str, rag_id: str, index: str) -> list: """Retrieve documents based on a query for a specific RAG ID and index. - Args: query (str): The query string to search for. rag_id (str): The ID of the RAG to search in. index (str): The index to search in. - Returns: list: A list of documents matching the query. """ - docs = await rag_factory.retrive_query(rag_name=rag_id, index=index, query=query) + docs = await rag_factory.retrieve_query(rag_name=rag_id, index=index, query=query) send_filter = [{"text": node.text, "score": node.score} for node in docs] return send_filter diff --git a/bolna/models.py b/bolna/models.py index 7bcc58aa..3c76b12b 100644 --- a/bolna/models.py +++ b/bolna/models.py @@ -171,6 +171,7 @@ class LLM(BaseModel): presence_penalty: Optional[float] = 0.0 provider: Optional[str] = "openai" base_url: Optional[str] = None + extra_config: Optional[ExtraConfig] = None class SIMPLE_LLM_AGENT(LLM): diff --git a/local_setup/ngrok-config.yml b/local_setup/ngrok-config.yml index 4907d3d9..2aedb97c 100644 --- a/local_setup/ngrok-config.yml +++ b/local_setup/ngrok-config.yml @@ -11,7 +11,7 @@ tunnels: bolna-app: addr: bolna-app:5001 proto: http - ingestion-app: - addr: ingestion-app:8000 - proto: http + # ingestion-app: + # addr: ingestion-app:8000 + # proto: http