API to load and query documents using RAG
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

origins = [
    'http://localhost:3000'
]

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def read_root():
    return {"Hello": "World"}

from routers.router import app as app_router
app.include_router(app_router)
main.py
Haste171/llamaindex-retrieval-api
import tempfile, os
from routers.utils.engine import Engine
from pathlib import Path
from llama_index import download_loader
from fastapi import APIRouter, UploadFile, HTTPException

engine = Engine()
router = APIRouter()

@router.post("/docx")
async def file(upload_file: UploadFile, namespace: str):
    """
    Loader: https://llamahub.ai/l/file-docx
    """
    file_preview_name, file_extension = os.path.splitext(upload_file.filename)
    if file_extension != '.docx':
        raise HTTPException(status_code=400, detail="File must be a docx")
    
    with tempfile.NamedTemporaryFile(delete=True, prefix=file_preview_name + '_', suffix=".docx") as temp_file:
        content = await upload_file.read()
        temp_file.write(content)
        DocxReader = download_loader("DocxReader")
        loader = DocxReader().load_data(file=Path(temp_file.name))
        engine.load(loader, namespace)
    
    return {'message': 'File uploaded successfully', 'filename': upload_file.filename, "namespace": namespace}
routers/loaders/docx.py
Haste171/llamaindex-retrieval-api
import tempfile, os
from routers.utils.engine import Engine
from pathlib import Path
from llama_index import download_loader
from fastapi import APIRouter, UploadFile, HTTPException

engine = Engine()
router = APIRouter()

@router.post("/ipynb")
async def file(upload_file: UploadFile, namespace: str):
    """
    Loader: https://llamahub.ai/l/file-ipynb
    """
    file_preview_name, file_extension = os.path.splitext(upload_file.filename)
    if file_extension != '.ipynb':
        raise HTTPException(status_code=400, detail="File must be a IPynb")
    
    with tempfile.NamedTemporaryFile(delete=True, prefix=file_preview_name + '_', suffix=".ipynb") as temp_file:
        content = await upload_file.read()
        temp_file.write(content)
        IPYNBReader = download_loader("IPYNBReader")
        loader = IPYNBReader(concatenate=True).load_data(file=Path(temp_file.name))
        engine.load(loader, namespace)
    
    return {'message': 'File uploaded successfully', 'filename': upload_file.filename, "namespace": namespace}
routers/loaders/ipynb.py
Haste171/llamaindex-retrieval-api
import tempfile, os
from routers.utils.engine import Engine
from llama_index import download_loader
from fastapi import APIRouter, UploadFile, HTTPException

engine = Engine()
router = APIRouter()

@router.post("/pdf")
async def file(upload_file: UploadFile, namespace: str):
    """
    Loader: https://llamahub.ai/l/file-pymu_pdf
    """
    file_preview_name, file_extension = os.path.splitext(upload_file.filename)
    if file_extension != '.pdf':
        raise HTTPException(status_code=400, detail="File must be a PDF")
    
    with tempfile.NamedTemporaryFile(delete=True, prefix=file_preview_name + '_', suffix=".pdf") as temp_file:
        content = await upload_file.read()
        temp_file.write(content)
        PyMuPDFReader = download_loader("PyMuPDFReader")
        loader = PyMuPDFReader().load(file_path=temp_file.name, metadata=True)
        engine.load(loader, namespace)
    
    return {'message': 'File uploaded successfully', 'filename': upload_file.filename, "namespace": namespace}
routers/loaders/pdf.py
Haste171/llamaindex-retrieval-api
from fastapi import APIRouter
from pydantic import BaseModel
from routers.utils.engine import Engine

query_engine = Engine()
router = APIRouter()

class QueryInput(BaseModel):
    prompt: str
    namespace: str
    temperature: float = 0.0

@router.post("/query")
async def query(input: QueryInput):
    response = query_engine.query(input.prompt, input.namespace, input.temperature)
    return {
        "input": input.prompt,
        "response": response
    }
routers/retrievers/query.py
Haste171/llamaindex-retrieval-api
from .retrievers import query
from .loaders import pdf, docx, ipynb
from fastapi import APIRouter

app = APIRouter()

app.include_router(query.router, tags=["retrieval"], prefix='/retrieval')
app.include_router(pdf.router, tags=["load"], prefix='/load')
app.include_router(docx.router, tags=["load"], prefix='/load')
app.include_router(ipynb.router, tags=["load"], prefix='/load')
routers/router.py
Haste171/llamaindex-retrieval-api
import pinecone
from llama_index.vector_stores import PineconeVectorStore
from llama_index.storage import StorageContext
from llama_index.llms import OpenAI
from llama_index.embeddings import OpenAIEmbedding
from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
from llama_index import ServiceContext
from decouple import config

class Engine:
    def __init__(self) -> None:
        self.openai_api_key = config('OPENAI_API_KEY')
        self.pinecone_api_key = config('PINECONE_API_KEY')
        self.pinecone_environment = config('PINECONE_ENVIRONMENT')
        self.pinecone_index = config('PINECONE_INDEX')
        self.embedding_model = OpenAIEmbedding(embed_batch_size=10, api_key=self.openai_api_key)


    def load(self, loader: object, namespace: str):
        """
        Designed to load or index data/documents from a given loader into a Pinecone vector store, using specific configurations and services, such as OpenAI's model for processing and embedding the data.
        """
        pinecone.init(api_key=self.pinecone_api_key, environment=self.pinecone_environment)
        pinecone_index = pinecone.Index(self.pinecone_index)
        vector_store = PineconeVectorStore(pinecone_index=pinecone_index, text_key="content", namespace=namespace)
        service_context = ServiceContext.from_defaults(llm=OpenAI(model="gpt-3.5-turbo-0613", api_key=self.openai_api_key), 
                                                       embed_model=self.embedding_model, chunk_size=1000)
        storage_context = StorageContext.from_defaults(vector_store=vector_store)
        VectorStoreIndex.from_documents(loader, service_context=service_context, storage_context=storage_context)


    def query(self, prompt: str, namespace: str, temperature: float):
        """
        Returns a response to the provided prompt based on the indexed data and the OpenAI model.
        Queries Pinecone vector store using OpenAI, and retrieve relevant responses based on the input prompt.
        """
        pinecone.init(api_key=self.pinecone_api_key, environment=self.pinecone_environment)
        pinecone_index = pinecone.Index(self.pinecone_index)
        service_context = ServiceContext.from_defaults(llm=OpenAI(model="gpt-3.5-turbo-0613", temperature=temperature, api_key=self.openai_api_key), 
                                                       embed_model=self.embedding_model, chunk_size=1000)
        vector_store = PineconeVectorStore(pinecone_index=pinecone_index, text_key="content", namespace=namespace)
        # retriever = VectorStoreIndex.from_vector_store(vector_store, service_context=service_context).as_retriever(similarity_top_k=1)
        index = VectorStoreIndex.from_vector_store(vector_store=vector_store, service_context=service_context)
        query_engine = index.as_query_engine(chat_mode="openai")
        response = query_engine.query(prompt)
        return response
    

    # def chat(self, prompt: str):
    #     service_context = ServiceContext.from_defaults(llm=OpenAI(model="gpt-3.5-turbo-0613", api_key=self.openai_api_key), 
    #                                                    embed_model=self.embedding_model, chunk_size=1000)
    #     data = SimpleDirectoryReader(input_dir="data").load_data()
    #     index = VectorStoreIndex.from_documents(data, service_context=service_context)
    #     chat_engine = index.as_chat_engine(chat_mode="openai", verbose=True, similarity_top_k=5, streaming=True)
    #     response = chat_engine.chat(prompt, function_call="query_engine_tool")
    #     return response
    

routers/utils/engine.py
Haste171/llamaindex-retrieval-api