|
|
from fastapi import FastAPI, Query |
|
|
from langchain_community.vectorstores import Chroma |
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings |
|
|
from together import Together |
|
|
import os |
|
|
|
|
|
embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2") |
|
|
vectordb = Chroma(persist_directory="chroma_db", embedding_function=embedding) |
|
|
|
|
|
TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY", "") |
|
|
client = Together(api_key=TOGETHER_API_KEY) |
|
|
|
|
|
def call_llama(prompt: str): |
|
|
response = client.chat.completions.create( |
|
|
model="meta-llama/Llama-3-8b-chat-hf", |
|
|
messages=[ |
|
|
{"role": "system", "content": "You are a helpful assistant."}, |
|
|
{"role": "user", "content": prompt} |
|
|
] |
|
|
) |
|
|
return response.choices[0].message.content |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
@app.get("/ask") |
|
|
async def ask(q: str = Query(..., description="Your question")): |
|
|
docs = vectordb.similarity_search(q, k=3) |
|
|
context = "\n".join([doc.page_content for doc in docs]) |
|
|
final_prompt = f"Use the context below to answer the question.\n\nContext:\n{context}\n\nQuestion: {q}" |
|
|
answer = call_llama(final_prompt) |
|
|
return {"answer": answer} |
|
|
|
