|
|
import os, json |
|
|
from langchain_community.vectorstores import FAISS |
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings |
|
|
from langchain.chains import RetrievalQAWithSourcesChain |
|
|
from langchain_community.chat_models import ChatOllama |
|
|
from langchain_community.llms import HuggingFaceHub |
|
|
from langchain.callbacks.base import BaseCallbackHandler |
|
|
from langchain_core.language_models.base import BaseLanguageModel |
|
|
import logging |
|
|
from langchain.globals import set_debug |
|
|
|
|
|
|
|
|
set_debug(True) |
|
|
_lc_logger = logging.getLogger("langchain") |
|
|
if not any(isinstance(h, logging.FileHandler) and getattr(h, "baseFilename", "").endswith("langchain_debug.jsonl") for h in _lc_logger.handlers): |
|
|
_fh = logging.FileHandler("langchain_debug.jsonl", mode="a", encoding="utf-8") |
|
|
_fh.setFormatter(logging.Formatter("%(message)s")) |
|
|
_lc_logger.addHandler(_fh) |
|
|
_lc_logger.setLevel(logging.DEBUG) |
|
|
|
|
|
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" |
|
|
|
|
|
|
|
|
def load_index(index_dir: str = "data"): |
|
|
embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL) |
|
|
store = FAISS.load_local(index_dir, embeddings, allow_dangerous_deserialization=True) |
|
|
with open(os.path.join(index_dir, "segments.json")) as f: |
|
|
segments = json.load(f) |
|
|
return store, segments |
|
|
|
|
|
|
|
|
class JSONLCallbackHandler(BaseCallbackHandler): |
|
|
"""Write simple LangChain events to a JSONL file so UI can display them.""" |
|
|
def __init__(self, path: str = "langchain_debug.jsonl"): |
|
|
self.path = path |
|
|
|
|
|
open(self.path, "w").close() |
|
|
|
|
|
def _write(self, record): |
|
|
import json, time |
|
|
record["ts"] = time.time() |
|
|
with open(self.path, "a", encoding="utf-8") as f: |
|
|
f.write(json.dumps(record) + "\n") |
|
|
|
|
|
def on_chain_start(self, serialized, inputs, **kwargs): |
|
|
self._write({"event": "chain_start", "name": serialized.get("name"), "inputs": inputs}) |
|
|
|
|
|
def on_chain_end(self, outputs, **kwargs): |
|
|
self._write({"event": "chain_end", "outputs": outputs}) |
|
|
|
|
|
def on_llm_start(self, serialized, prompts, **kwargs): |
|
|
self._write({"event": "llm_start", "prompts": prompts}) |
|
|
|
|
|
def on_llm_end(self, response, **kwargs): |
|
|
self._write({"event": "llm_end", "response": str(response)}) |
|
|
|
|
|
def on_retriever_end(self, documents, **kwargs): |
|
|
from langchain.docstore.document import Document |
|
|
preview = [doc.page_content[:200] if isinstance(doc, Document) else str(doc) for doc in documents] |
|
|
self._write({"event": "retriever_end", "documents": preview}) |
|
|
|
|
|
|
|
|
def get_model(model_name: str, hf_token: str = None, callbacks: list = None) -> BaseLanguageModel: |
|
|
"""Return a model instance based on the model name. |
|
|
|
|
|
Args: |
|
|
model_name: Name of the model to use |
|
|
hf_token: Hugging Face API token (required for flan-t5-base) |
|
|
callbacks: List of callbacks to use |
|
|
""" |
|
|
if model_name == "flan-t5-base": |
|
|
if not hf_token: |
|
|
raise ValueError( |
|
|
"Hugging Face API token is required for flan-t5-base. " |
|
|
"Please provide your Hugging Face token in the UI or use a local model." |
|
|
) |
|
|
return HuggingFaceHub( |
|
|
repo_id="google/flan-t5-base", |
|
|
huggingfacehub_api_token=hf_token, |
|
|
model_kwargs={"temperature": 0.1, "max_length": 512}, |
|
|
callbacks=callbacks |
|
|
) |
|
|
else: |
|
|
return ChatOllama(model=model_name, callbacks=callbacks) |
|
|
|
|
|
|
|
|
def build_chain(store, model_name: str = "phi3", hf_token: str = None): |
|
|
"""Return a RetrievalQA chain using the specified model. |
|
|
|
|
|
Args: |
|
|
store: Vector store with document embeddings |
|
|
model_name: Name of the model to use |
|
|
hf_token: Hugging Face API token (required for flan-t5-base) |
|
|
""" |
|
|
callback = JSONLCallbackHandler() |
|
|
llm = get_model(model_name, hf_token, [callback]) |
|
|
return RetrievalQAWithSourcesChain.from_chain_type( |
|
|
llm=llm, |
|
|
retriever=store.as_retriever(k=4, callbacks=[callback]), |
|
|
return_source_documents=True, |
|
|
verbose=True, |
|
|
) |
|
|
|
