Practical notes: Building a Company AI Assistant with RAG, Agents, and LangGraph

The following notes reconstruct a practical path around “Building a Company AI Assistant with RAG, Agents, and LangGraph”. Emphasis stays on contracts, checks, and drop-in code placeholders rather than motivational framing. When working through Overview, write down the contract first: required inputs, success signal, and what happens on partial failure. That checklist keeps later code changes honest. Prefer small, testable units over sprawling scripts. When a step fails, the failure should point at a single responsibility rather than a tangled pipeline.

rag-agent-project/
│
├── documents/
├── faiss_index/
│
├── .env
├── .gitignore
├── requirements.txt
│
├── config.py
├── ingest.py
├── rag.py
├── query_rewriter.py
├── reranker.py
├── tools.py
├── graph.py
├── main.py
└── app.py

pip install -r requirements.txt

.env

OPENAI_API_KEY=your_openai_api_key

import os
from dotenv import load_dotenv

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
    raise ValueError("OPEN_API_KEY is not set."
                     "Add it to yor .env file.")

LLM_MODEL = "gpt-4.1-mini"

EMBEDDING_MODEL = "text-embedding-3-small"

HUGGINGFACE_EMBEDDING = "sentence-transformers/all-MiniLM-L6-v2"
FAISS_PATH = "faiss_index"

DATABASE_PATH = "company_assistant.db"

documents/
├── employee_handbook.pdf
├── hr_policy.pdf
└── it_policy.pdf

from pathlib import Path

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

from config import HUGGINGFACE_EMBEDDING , FAISS_PATH


DOCUMENTS_DIR = Path("documents")


def load_documents():

    all_documents = []

    pdf_files = list(DOCUMENTS_DIR.glob("*.pdf"))

    if not pdf_files:
        raise FileNotFoundError(
            "No pdf file found in documents/"
        )

    for pdf_file in pdf_files:

        print(f"Loading: {pdf_file}")

        loader = PyPDFLoader(str(pdf_file))
        documents = loader.load()
        for document in documents:
            document.metadata["source"] = pdf_file.name
            if "handbook" in pdf_file.name.lower():
                document.metadata["document_type"] = "handbook"

            elif "policy" in pdf_file.name.lower():

                document.metadata["document_type"] = "policy"

            else:

                document.metadata["document_type"] = "other"

            all_documents.append(document)

    return all_documents


def main():
    # 1. Load PDFs
    documents = load_documents()

    print(f"\nLoaded {len(documents)} pages")

    # 2. Split documents
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100
    )

    chunks = text_splitter.split_documents(documents)

    print(f"Split into {len(chunks)} chunks.")

    # 3. Create embeddings
    embeddings = HuggingFaceEmbeddings(
        model_name=HUGGINGFACE_EMBEDDING
    )

    # 4. Create FAISS database
    vector_store = FAISS.from_documents(
        chunks,
        embeddings
    )


    # 5. Save FAISS database
    vector_store.save_local(FAISS_PATH)

    print(
        f"\nFAISS index saved to: {FAISS_PATH}"
    )


if __name__ == "__main__":
    main()

Loading: documents/employee_handbook.pdf
Loading: documents/hr_policy.pdf
Loading: documents/it_policy.pdf
Loaded 20 pages.
Created 60 chunks.
FAISS index created successfully!

Chunk 1
-------------------------
A B C D E F G H I J
              |
              | overlap
              v
Chunk 2
              E F G H I J K L M N

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from config import(
    HUGGINGFACE_EMBEDDING,
    FAISS_PATH
)

embeddings = HuggingFaceEmbeddings(
        model_name=HUGGINGFACE_EMBEDDING
    )

vector_store = FAISS.load_local(
    FAISS_PATH,
    embeddings,
    allow_dangerous_deserialization=True
)

from rag import vector_store

from langchain_huggingface import (HuggingFaceEmbeddings,HuggingFacePipeline)
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-0.5B-Instruct",
    max_new_tokens=256,
    temperature=0.1
)

llm = HuggingFacePipeline(
    pipeline=pipe
)

def rewrite_query(question:str)-> str:
    prompt = f"""
Rewrite the user's question into a concise search query for company document database.

The query should contain important keywords that are likely to appear in company documents.

Do not answer the question.

User question:{question}

search query:
"""

    response = llm.invoke(prompt)

    return response

Chunk 1
Chunk 2
Chunk 3
Chunk 4
Chunk 5

from langchain_huggingface import (HuggingFaceEmbeddings,HuggingFacePipeline)
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-0.5B-Instruct",
    max_new_tokens=256,
    temperature=0.1
)

llm = HuggingFacePipeline(
    pipeline=pipe
)

def rerank_documents(query:str, documents:list,top_n: int = 3):

    scored_documents =[]

    for document in documents:
        prompt = f"""
You are evaluating document releveance.

Rate how relevant the document is to the user's question.

Return ONLY an intefer from 0 to 10.

0 = competely irrelevent
10 = extreamly relevent

Question: {query}

Document:{document.page_content}

"""
        response = llm.invoke(prompt)
        try:
            score = int(response.content.strip())
        except ValueError:
            score = 0

        scored_documents.append((score,document))

        #Highest score first
        scored_documents.sort(
            key=lambda item:item[0],
            reverse=True
        )

        return[
            document
            for score, document
            in scored_documents[:top_n]
        ]

from langchain.tools import tool
from rag import vector_store
from query_rewriter import(rewrite_query)

from reranker import (rerank_documents)


@tool
def company_document_search(query: str) -> str:
    """
    Search company PDF documents for information about

    company policies, employee benifits, leave, holidays, work from home,
    IT rules, working hours, and othher internal information.

    """
    #1. Rewrite Query
    imporved_query = rewrite_query(query)
    print(f"\n[Rewritten Query]"
          f"\n{imporoved_query}")

    #2.  Vector serach
    documents =  (vector_store.similarity_search(
        imporved_query,
        k=5
    ))

    if not documents:
        return(
            "No relevent information"
            "was found in the company documents"
        )

    #Reranking
    documents = rerank_documents(
        imporved_query, documents,top_n=3
    )

    #4.Formate context
    result = []
    for document in documents:
        source = document.metadata.get("source","Unknown")
        page = document.metadata.get("page","Unknown")

        #PyPDFLoader user zero-based pages
        if isinstance(page, int):
            page = page+1

        results.append(f"""
SOURCE:{source}
PAGE:{page}

CONTENT: {document.page_content}
""")
        return "\n".join(result)


@tool
def calculator(expression: str) -> str:
    """
    Calculate a mathematical expression.
    """
    try:

       # result = eval(expression)
       # for learning only
       # Do not use eval() in production

        allowed = ("0123456789""+-*/(). ")
        if not all(char in allowed for char in expression):
            return ("Invalid mathmatical expression")

        result = eval(expression,
                      {"__builtins__":{}},
                      {}
                      )

        return str(result)
    except Exception:
        return "Unable to calculate the expression."

from langchain_ollama import ChatOllama

from langgraph.graph import (
    StateGraph,
    MessagesState,
    START
)

from langgraph.prebuilt import (
    ToolNode,
    tools_condition
)

from langgraph.checkpoint.memory import InMemorySaver

from tools import (
    company_document_search,
    calculator
)

# Local Qwen model
llm = ChatOllama(
    model="qwen2.5:3b",
    temperature=0
)

# Tools
tools = [
    company_document_search,
    calculator
]


llm_with_tools = llm.bind_tools(tools)

# Agent

def agent_node(state: MessagesState):

    response = llm_with_tools.invoke(
        state["messages"]
    )

    return {
        "messages": [response]
    }

# Graph
graph_builder = StateGraph(MessagesState)


graph_builder.add_node(
    "agent",
    agent_node
)


graph_builder.add_node(
    "tools",
    ToolNode(tools)
)

# START → Agent
graph_builder.add_edge(
    START,
    "agent"
)

# Agent → Tool OR END
graph_builder.add_conditional_edges(
    "agent",
    tools_condition
)

# Tool → Agent
graph_builder.add_edge(
    "tools",
    "agent"
)

# Checkpointer
checkpointer = InMemorySaver()

# Compile

graph = graph_builder.compile(
    checkpointer=checkpointer
)

SOURCE: hr_policy.pdf
PAGE: 4

Source: hr_policy.pdf, page 4.

from langchain_core.messages import HumanMessage, SystemMessage
from graph import graph

SYSTEM_PROMPT="""
You are an intelligent company assistant.

You answer question using company documents and avalable tools.

Rules:
1. Use company_document_search whenever the user asks about company-specific information.

2. Never invent company policies,  benifits, rules, or  procedures.

3. Use the calculator when mathmatical calculation are required.

4. When company documents are used, mention the source document and page.

5. If the company documents do not contain the requested information, clearly say
that you could not find it.

6. Use previous conversation context to undersatnd follow-up questions.

7. Give concise and clear answer
"""

def chat(
        question:str,
        thread_id: str):
    config = {
        "configurable":{"thread_id":thread_id}
    }

    #Check whether thhis conversation already exists
    state = graph.get_state(config)

    if not state.values:
        messages = [
            SystemMessage(content=SYSTEM_PROMPT),
            HumanMessage(content=question)
        ]
    else:
        messages=[
            HumanMessage(content=question)]

    result = graph.invoke({
        "messages":messages},
        config=config
    )
    return result["messages"][-1].content

def main():
    print("===========================================")
    print("CoOMPANY AI ASSISTANT")
    print("===========================================")
    print("\nType 'exist' to quit.")
    thread_id = input("\nEnter conversation ID:")

    while True:
        question = input("\nYou: ")

        if question.lower() in [
            "exit",
            "quit"
        ]:
            break

        try:
            answer = chat(question,thread_id)
            print(f"\nAI: {answer}")

        except Exception as e:
            print(f"\nError: {e}")

if __name__ == "__main__":
    main()

import streamlit as st
from langchain_core.messages import (
    HumanMessage,
    AIMessage,
    SystemMessage
)

from graph import graph

# Page configuration
st.set_page_config(
    page_title="Company AI Assistant",
    page_icon="🤖",
    layout="centered"
)

# System prompt
SYSTEM_PROMPT = """
You are an intelligent company assistant.

You answer questions using company documents
and available tools.

Rules:

1. Use company_document_search whenever
   the user asks about company-specific
   information.

2. Never invent company policies,
   benefits, rules, or procedures.

3. Use the calculator when mathematical
   calculations are required.

4. When company documents are used,
   mention the source document and page.

5. If the company documents do not contain
   the requested information, clearly say
   that you could not find it.

6. Use previous conversation context to
   understand follow-up questions.

7. Give concise and clear answers.
"""

# Sidebar
st.sidebar.title("⚙️ Settings")

thread_id = st.sidebar.text_input(
    "Conversation ID",
    value="streamlit-user"
)

st.sidebar.markdown(
    """
### About

This application demonstrates:

- RAG
- AI Agents
- LangGraph
- FAISS
- Query rewriting
- Reranking
- Conversation memory
- Tool calling
"""
)

# Main UI
st.title("🏢 Company AI Assistant")

st.write(
    "Ask questions about company policies, "
    "benefits, leave, IT rules, and more."
)

# Chat history in Streamlit

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Display previous messages
for message in st.session_state.chat_history:

    with st.chat_message(
        message["role"]
    ):

        st.markdown(
            message["content"]
        )


# User input
question = st.chat_input(
    "Ask something about the company..."
)


if question:
    # Display user question
    with st.chat_message("user"):

        st.markdown(question)


    st.session_state.chat_history.append(
        {
            "role": "user",
            "content": question
        }
    )

    # Build messages
    config = {
        "configurable": {
            "thread_id": thread_id
        }
    }

    # Check whether this thread
    # already has state
    state = graph.get_state(
        config
    )
    if not state.values:

        messages = [
            SystemMessage(
                content=SYSTEM_PROMPT
            ),
            HumanMessage(
                content=question
            )
        ]

    else:
        messages = [
            HumanMessage(
                content=question
            )
        ]

    # Call LangGraph
    with st.chat_message("assistant"):
        with st.spinner(
            "Thinking..."
        ):
            try:
                result = graph.invoke(
                    {
                        "messages": messages
                    },
                    config=config
                )
                answer = result[
                    "messages"
                ][-1].content
                st.markdown(answer)

                # Save answer
                st.session_state.chat_history.append(
                    {
                        "role": "assistant",
                        "content": answer
                    }
                )
            except Exception as e:

                st.error(
                    f"Something went wrong: {e}"
                )

Operational checklist

For Operational checklist, define the inputs, the owner of the step, and the exit criteria before changing code. Operators should be able to re-run the step from a known checkpoint without guessing hidden state.

Document the happy path and the recovery path together. Retries, human gates, and dead-letter handling are part of the product, not later polish.

Cite the passages that actually grounded the answer. Without citations, operators cannot tell hallucination from an indexing gap.

Checkpoint after expensive steps. Resume should not re-bill the same LLM call when an operator retries a later node.

Pin dependency versions and record the image digest that ran the demo. Reproducibility beats tribal knowledge.

Record timings and token or query cost next to functional results. Cost visibility early prevents surprise bills when the path moves from demo to shared environments.

Before promoting the stack, freeze versions, capture a golden transcript for the critical path, and confirm rollback steps. Shared environments need rate limits, tenancy checks, and a clear owner for secret rotation. Prefer boring reliability over clever one-off demos.

Batch note for 656e1cd26886: keep provider keys out of the repo, set a per-session token ceiling, and store transcripts next to the eval fixtures so later model swaps stay comparable.

For hardening note 0, define the inputs, the owner of the step, and the exit criteria before changing code. Operators should be able to re-run the step from a known checkpoint without guessing hidden state. Keep configuration outside application code. Environment files, secret stores, and feature flags belong in one place operators can audit without reading the whole graph.

Hardening detail 0/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

When working through hardening note 1, write down the contract first: required inputs, success signal, and what happens on partial failure. That checklist keeps later code changes honest. Prefer small, testable units over sprawling scripts. When a step fails, the failure should point at a single responsibility rather than a tangled pipeline.

Hardening detail 1/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

hardening note 2 works best when treated as a measurable surface. Capture one golden transcript, one failure case, and the rollback note before expanding scope. Record timings and token or query cost next to functional results. Cost visibility early prevents surprise bills when the path moves from demo to shared environments.

Hardening detail 2/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

For hardening note 3, define the inputs, the owner of the step, and the exit criteria before changing code. Operators should be able to re-run the step from a known checkpoint without guessing hidden state. Document the happy path and the recovery path together. Retries, human gates, and dead-letter handling are part of the product, not later polish.

Hardening detail 3/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

When working through hardening note 4, write down the contract first: required inputs, success signal, and what happens on partial failure. That checklist keeps later code changes honest. Treat this stage as a contract between inputs and validated outputs. Name the artifacts, define success checks, and refuse silent partial completion.

Hardening detail 4/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

hardening note 5 works best when treated as a measurable surface. Capture one golden transcript, one failure case, and the rollback note before expanding scope. Keep configuration outside application code. Environment files, secret stores, and feature flags belong in one place operators can audit without reading the whole graph.

Hardening detail 5/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

For hardening note 6, define the inputs, the owner of the step, and the exit criteria before changing code. Operators should be able to re-run the step from a known checkpoint without guessing hidden state. Prefer small, testable units over sprawling scripts. When a step fails, the failure should point at a single responsibility rather than a tangled pipeline.

Hardening detail 6/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

When working through hardening note 7, write down the contract first: required inputs, success signal, and what happens on partial failure. That checklist keeps later code changes honest. Record timings and token or query cost next to functional results. Cost visibility early prevents surprise bills when the path moves from demo to shared environments.

Hardening detail 7/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

hardening note 8 works best when treated as a measurable surface. Capture one golden transcript, one failure case, and the rollback note before expanding scope. Document the happy path and the recovery path together. Retries, human gates, and dead-letter handling are part of the product, not later polish.

Hardening detail 8/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.

For hardening note 9, define the inputs, the owner of the step, and the exit criteria before changing code. Operators should be able to re-run the step from a known checkpoint without guessing hidden state. Treat this stage as a contract between inputs and validated outputs. Name the artifacts, define success checks, and refuse silent partial completion.

Hardening detail 9/780: measure wall time, error class, and token spend for this note, then decide whether to keep the change based on a fixed question set rather than anecdote.