File size: 2,430 Bytes
14faba3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import numpy as np
import os
os.environ['NUMPY_IMPORT'] = 'done' # This ensures numpy is loaded
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import CharacterTextSplitter
from app.config import CHROMA_DB_DIR
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from dotenv import load_dotenv
load_dotenv()
OPENAI_ROUTER_TOKEN=os.getenv("OPENROUTER")
# Embeddings
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
# Chroma DB
db = Chroma(persist_directory=CHROMA_DB_DIR, embedding_function=embeddings)
from langchain.docstore.document import Document
def add_document(file_path: str, user_id: str):
# Load file
if file_path.lower().endswith(".pdf"):
loader = PyPDFLoader(file_path)
elif file_path.lower().endswith(".txt"):
loader = TextLoader(file_path, encoding="utf-8")
else:
raise RuntimeError(f"Unsupported file type: {file_path}")
documents = loader.load()
# Split into chunks
splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
docs = splitter.split_documents(documents)
# Add metadata directly to Document objects
docs_with_metadata = [
Document(page_content=d.page_content, metadata={"user_id": user_id, "filename": os.path.basename(file_path)})
for d in docs
]
# Add to vector store
db.add_documents(docs_with_metadata)
def get_qa_chain(user_id: str):
"""
Return a RetrievalQA pipeline for a specific user using OpenRouter's Phi-3 Medium Instruct model.
Args:
user_id (str): Unique identifier for the user.
"""
# Initialize LLM with OpenRouter
llm = ChatOpenAI(
openai_api_key=OPENAI_ROUTER_TOKEN, # your OpenRouter API key
model="meta-llama/llama-4-scout:free", # free OpenRouter model
temperature=0,
max_tokens=512,
openai_api_base="https://openrouter.ai/api/v1" # OpenRouter endpoint
)
# Create retriever filtered by user_id
retriever = db.as_retriever(search_kwargs={"filter": {"user_id": user_id}})
# Build RetrievalQA pipeline
qa = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, chain_type="stuff")
return qa
|