diff --git a/README.md b/README.md index ddca3534..50e8a4af 100644 --- a/README.md +++ b/README.md @@ -203,6 +203,30 @@ input_dir = "/local/documents/path" documents = local_document_provider(input_dir=input_dir) ``` +### FastAPI Integration in 1-Line 🚀 + +Creating a FastAPI application integrated with AutoLLM has never been easier. Follow the quick guide below to get started. + +#### Create Your FastAPI Application + +In your `main.py`, include the following line of code: + +```python +from autollm import create_web_app + +app = create_web_app(config, env) +``` + +Here, `config` and `env` should be replaced by your configuration and environment file paths. + +#### Run Your Application + +After creating your FastAPI app, run the following command in your terminal to get it up and running: + +```bash +uvicorn main:app +``` + ______________________________________________________________________ ## FAQ diff --git a/autollm/app/main.py b/autollm/app/main.py index c3a8cb8f..6c7de0ea 100644 --- a/autollm/app/main.py +++ b/autollm/app/main.py @@ -1,6 +1,7 @@ import logging from fastapi import FastAPI, HTTPException, Query +from pydantic import BaseModel, Field from autollm.app.docs import description, openapi_url, tags_metadata, terms_of_service, title, version from autollm.app.utils import load_config_and_initialize_engines @@ -8,6 +9,11 @@ logging.basicConfig(level=logging.INFO) +class QueryPayload(BaseModel): + task: str = Field(..., description="Task to execute") + user_query: str = Field(..., description="User's query") + + # Function to create the FastAPI web app def create_web_app(config_file_path: str, env_file_path: str = None): app = FastAPI( @@ -22,9 +28,10 @@ def create_web_app(config_file_path: str, env_file_path: str = None): query_engines = load_config_and_initialize_engines(config_file_path, env_file_path) @app.post("/query") - async def query( - task: str = Query(..., description="Task to execute"), - user_query: str = Query(..., description="User's query")): + async def query(payload: QueryPayload): + task = payload.task + user_query = payload.user_query + if task not in query_engines: raise HTTPException(status_code=400, detail="Invalid task name") diff --git a/autollm/auto/vector_store_index.py b/autollm/auto/vector_store_index.py index 48f5cc78..b3681d52 100644 --- a/autollm/auto/vector_store_index.py +++ b/autollm/auto/vector_store_index.py @@ -40,7 +40,7 @@ def from_defaults( index (VectorStoreIndex): The initialized Vector Store index instance for given vector store type and parameter set. """ if documents is None: - documents = [Document.example()] + documents = Document.example() if vector_store_type == "VectorStoreIndex": index = VectorStoreIndex.from_documents(documents=[documents], *args, **kwargs) else: diff --git a/tests/test_app.py b/tests/test_app.py new file mode 100644 index 00000000..854b3c1a --- /dev/null +++ b/tests/test_app.py @@ -0,0 +1,61 @@ +# Importing the necessary modules for testing and mocking +from unittest.mock import mock_open, patch + +from fastapi.testclient import TestClient +from llama_index.indices.query.base import BaseQueryEngine + +from autollm.app.main import create_web_app, load_config_and_initialize_engines + +# Mock file content that mimics a YAML configuration file +mock_file_content = "task1: {system_prompt: 'prompt1'}" + + +# 1. Testing Configuration Loading and Initialization +@patch('builtins.open', new_callable=mock_open, read_data=mock_file_content) +@patch('yaml.safe_load') +def test_load_config_and_initialize_engines(mock_safe_load, mock_file): + # Mock the yaml.safe_load to return a sample config + sample_config = { + 'summarize': { + 'system_prompt': 'You are a friendly chatbot that can summarize documents.' + }, + 'qa': { + 'system_prompt': 'You are a friendly chatbot that can answer questions.' + } + } + mock_safe_load.return_value = sample_config + + query_engines = load_config_and_initialize_engines('mock_path') + + # Assert that the open function was called with the correct file path + mock_file.assert_any_call('mock_path') + + # Validate the type and configuration of each query engine + for task in sample_config.keys(): + assert isinstance(query_engines[task], BaseQueryEngine) + + +# 2. Testing FastAPI App Creation +@patch('builtins.open', new_callable=mock_open, read_data=mock_file_content) +def test_create_web_app(mock_file): + # Assuming the function create_web_app is being imported correctly + app = create_web_app('mock_path') + + # Validate presence of the query endpoint + assert any(route.path == "/query" for route in app.routes) + + +# 3. Testing Query Endpoint +@patch('builtins.open', new_callable=mock_open, read_data=mock_file_content) +def test_query_endpoint(mock_file): + # Assuming the function create_web_app is being imported correctly + app = create_web_app('mock_path') + client = TestClient(app) + + # Test with a valid task + response = client.post("/query", json={"task": "task1", "user_query": "test query"}) + assert response.status_code == 200 + + # Test with an invalid task + response = client.post("/query", json={"task": "invalid_task", "user_query": "test query"}) + assert response.status_code == 400