deepseek_langgraph.py

import os
import torch
from langgraph.graph import StateGraph
from transformers import AutoModelForCausalLM, AutoTokenizer
from sentence_transformers import SentenceTransformer
import faiss

MODEL_NAME = os.getenv("DEEPSEEK_MODEL", "deepseek-ai/deepseek-llm-67b-chat")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    torch_dtype=torch.float16,
    device_map="auto",
)

# Lightweight in‑memory vector store for demo purposes
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
index = faiss.IndexFlatIP(384)
docs = []

def add_doc(text: str) -> None:
    vec = embedder.encode([text])
    index.add(vec)
    docs.append(text)

def retrieve(query: str, k: int = 3):
    if len(docs) == 0:
        return []
    qvec = embedder.encode([query])
    _, I = index.search(qvec, k)
    return [docs[i] for i in I[0] if i < len(docs)]


def node_user(state):
    return state

def node_retrieve(state):
    state["context"] = retrieve(state["prompt"])
    return state

def node_llm(state):
    context = "\n".join(state["context"])
    prompt = (
        "Answer the question using the context below.\n"
        "Context:\n" + context + "\n\nQuestion: " + state["prompt"] + "\nAnswer:"
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=512)
    state["answer"] = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return state

# Build LangGraph
lg = StateGraph()
lg.add_node("user", node_user)
lg.add_node("retrieve", node_retrieve)
lg.add_node("deepseek", node_llm)
lg.add_edge("user", "retrieve")
lg.add_edge("retrieve", "deepseek")
lg.set_entrypoint("user")
lg.set_exit("deepseek")

pipeline = lg.compile()

if __name__ == "__main__":
    add_doc("LangGraph is a Python library for building stateful graphs around LLM calls and tool use.")
    out = pipeline({"prompt": "What is LangGraph?"})
    print(out["answer"])