Changes from ggruber193/polars-docu-chat-rag

app.py

@@ -0,0 +1,84 @@
+import gradio as gr
+from uuid import uuid4
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.store.memory import InMemoryStore
+
+from src.rag_lanchain import graph_builder
+
+
+memory = MemorySaver()
+in_memory_store = InMemoryStore()
+graph = graph_builder.compile(checkpointer=memory, store=in_memory_store)
+
+
+def respond(msg, config):
+    role_dict = {"ai": "assistant", "human": "user"}
+    if len(msg) == 0:
+        gr.Warning("Chat messages cannot be empty")
+        history = []
+        for hist in graph.get_state_history(config):
+            history = [{"role": role_dict.get(i.type, i.type), "content": i.content} for i in
+                       hist.values["messages"]]
+            break
+        return "", history
+    events = graph.stream(
+        {"messages": [{"role": "user", "content": msg}]},
+        config,
+        stream_mode="values",
+    )
+    events = list(events)
+    conversation = events[-1]["messages"]
+    conversation = [{"role": role_dict.get(i.type, i.type), "content": i.content} for i in conversation]
+    return "", conversation
+
+
+def init_chat_state():
+    return {"configurable": {"thread_id": str(uuid4()).replace('-', '_')}}
+
+
+css = """
+.centered-container {
+    max-width: 1000px;
+    margin: 0 auto;
+}
+"""
+
+THEME = gr.themes.Ocean()
+
+demo = gr.Blocks(theme=THEME, fill_width=False, fill_height=True, css=css)
+
+with demo:
+    config_state = gr.State(init_chat_state)
+    with gr.Column(elem_classes="centered-container"):
+        gr.Markdown("""
+        # 💬 Polars Python Chatbot
+        ### Ask anything about the [Polars](https://pola-rs.github.io/polars/) Python package!  
+        ### This chatbot uses a database of embeddings generated from the official documentation to help you find accurate and relevant answers about using Polars for data manipulation in Python.
+        """)
+
+        chatbot = gr.Chatbot(
+            label=None,
+            type="messages",
+            show_label=False,
+            height=400,
+        )
+
+        with gr.Row(equal_height=True):
+            msg = gr.Textbox(
+                placeholder="Type your message here...",
+                show_label=False,
+                lines=3,
+                max_lines=3,
+                scale=5,
+            )
+            send_btn = gr.Button("Send", variant="primary", scale=1)
+
+        with gr.Row():
+            clear = gr.ClearButton([msg, chatbot], value="Clear Chat", variant="secondary")
+
+    send_btn.click(respond, [msg, config_state], [msg, chatbot])
+    msg.submit(respond, [msg, config_state], [msg, chatbot])
+
+if __name__ == '__main__':
+    demo.launch()

requirements.txt

@@ -0,0 +1,10 @@
+beautifulsoup4~=4.13.4
+markdown~=3.8
+langchain~=0.3.23
+transformers~=4.51.3
+torch~=2.6.0
+tqdm~=4.67.1
+qdrant_client
+langgraph~=0.3.31
+gradio~=5.25.2
+langchain_google_genai~=2.1.3

src/config.py

@@ -0,0 +1,19 @@
+import os
+
+
+EMBEDDING_MODEL = "thenlper/gte-small"
+
+QDRANT_COLLECTION_NAME = "polars-documentation"
+QDRANT_URL = os.environ.get("QDRANT_URL", "")
+QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", "")
+CHAT_API_KEY = os.environ.get("CHAT_API_KEY", "")
+
+
+def get_qdrant_config():
+    from qdrant_client import models
+    QDRANT_COLLECTION_CONFIG = {
+        "collection_name": QDRANT_COLLECTION_NAME,
+        "vectors_config": models.VectorParams(size=384, distance=models.Distance.COSINE),   # on_disk=True),
+        # "hnsw_config": models.HnswConfigDiff(on_disk=True)
+    }
+    return QDRANT_COLLECTION_CONFIG

src/data_processing/process_markdown.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from bs4 import BeautifulSoup
+from langchain_core.documents import Document
+from markdown import markdown
+from pathlib import Path
+from langchain.text_splitter import MarkdownTextSplitter, MarkdownHeaderTextSplitter, TextSplitter
+
+from src.utils import batched
+
+
+def read_markdown_file(path: str | Path) -> [str, str]:
+    path = Path(path)
+    with open(path, 'r', encoding="utf8") as f_r:
+        text = f_r.read()
+
+    # text = markdown(text)
+    # text = ''.join(BeautifulSoup(text).findAll(text=True))
+    return text, str(path)
+
+
+def split_markdown(md: str | list[str],
+                   metadata=dict[str, Any] | list[dict[str, Any]],
+                   chunk_size=512,
+                   overlap=64,
+                   splitter: TextSplitter = None) -> list[Document]:
+    if isinstance(md, str):
+        md = [md]
+        if isinstance(metadata, list):
+            raise ValueError("metadata should be a single dict")
+        metadata = [metadata]
+    if splitter is None:
+        headers_to_split_on = [
+            ("#", "Header 1"),
+            ("##", "Header 2"),
+            ("###", "Header 3"),
+        ]
+        md = [MarkdownHeaderTextSplitter(headers_to_split_on, strip_headers=False).split_text(i) for i in md]
+        metadata = [{**metadata[i], **text.metadata} for i, text_split in enumerate(md) for text in text_split]
+        md = [j.page_content for i in md for j in i]
+        splitter = MarkdownTextSplitter(chunk_size=chunk_size, chunk_overlap=overlap)
+
+    docs = splitter.create_documents(md, metadata)
+    return docs
+
+
+def process_markdown_files(paths: list[str | Path], batch_size=1, chunk_size=512, overlap=64):
+    for files in batched(paths, batch_size):
+        mds_w_paths = [read_markdown_file(i) for i in files]
+        metadata = [{"path": md_path} for _, md_path in mds_w_paths]
+        md = [md for md, _ in mds_w_paths]
+        docs = split_markdown(md, metadata, chunk_size=chunk_size, overlap=overlap)
+        yield [i.page_content for i in docs], [i.metadata for i in docs]

src/database/qdrant_store.py

@@ -0,0 +1,83 @@
+from typing import Any
+
+from qdrant_client import QdrantClient, models
+from uuid import uuid4
+from transformers import PreTrainedModel
+
+from src.config import QDRANT_COLLECTION_NAME, QDRANT_URL, QDRANT_API_KEY, EMBEDDING_MODEL
+from src.embeddings import TextEmbedder
+
+
+class QdrantStore:
+    def __init__(self, client: QdrantClient, collection_config=None):
+        self.client = client
+        self.collection_names = set([i.name for i in client.get_collections().collections])
+
+        if collection_config is not None:
+            self.create_collection(collection_config)
+
+    def create_collection(self, collection_config: dict):
+        collection_name = collection_config["collection_name"]
+        if not self.client.collection_exists(collection_name):
+            self.client.create_collection(**collection_config)
+            self.collection_names.add(collection_name)
+
+    def _check_collection_name(self, collection_name):
+        if collection_name not in self.collection_names:
+            raise ValueError(f"Collection: {collection_name} does not exist.")
+
+    def upsert_points(self,
+                      vectors: Any | list[Any],
+                      payloads: dict | list[dict],
+                      collection_name: str):
+        self._check_collection_name(collection_name)
+
+        ids = [str(uuid4()) for _ in payloads]
+
+        self.client.upsert(
+            collection_name=collection_name,
+            points=models.Batch(
+                ids=ids,
+                payloads=payloads,
+                vectors=vectors
+            )
+        )
+
+    def delete_points(self,
+                      filters: dict[str, list[models.FieldCondition]],
+                      collection_name: str):
+        self._check_collection_name(collection_name)
+
+        self.client.delete(
+            collection_name=collection_name,
+            points_selector=models.Filter(**filters)
+        )
+
+    def delete_points_by_match(self,
+                               key_value: tuple[str, list[str] | str],
+                               collection_name: str):
+        key, values = key_value
+        if isinstance(values, str):
+            values = [values]
+        filter = {"must": [models.FieldCondition(key=key, match=models.MatchAny(any=values))]}
+        self.delete_points(filter, collection_name)
+
+    def get_topk_points_single(self,
+                               query: Any | str,
+                               collection_name: str,
+                               k=5):
+        responses = self.client.query_points(collection_name=collection_name,
+                                             query=query,
+                                             limit=k)
+
+        return [i.payload["text"] for i in responses.points]
+
+
+if __name__ == '__main__':
+    client = QdrantClient(QDRANT_URL, api_key=QDRANT_API_KEY)
+    qdrant_store = QdrantStore(client)
+    embedding_model = TextEmbedder(modelname=EMBEDDING_MODEL)
+    query = "How to filter a dataframe"
+    query_emb = embedding_model.embed_text(query)
+    responses = qdrant_store.get_topk_points_single(query_emb[0], collection_name=QDRANT_COLLECTION_NAME)
+    print(responses)

src/embeddings.py

@@ -0,0 +1,37 @@
+from transformers import AutoModel, AutoTokenizer
+from torch import Tensor
+from torch import functional as F
+
+from src.config import EMBEDDING_MODEL
+from src.utils import batched
+
+
+class TextEmbedder:
+    def __init__(self, modelname=EMBEDDING_MODEL, max_length=512):
+        self.tokenizer = AutoTokenizer.from_pretrained(modelname)
+        self.model = AutoModel.from_pretrained(modelname)
+        self.max_length = max_length
+
+    @staticmethod
+    def average_pool(last_hidden_states: Tensor,
+                     attention_mask: Tensor) -> Tensor:
+        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+    def embed_text(self, text: str | list[str], batch_size=128):
+        if isinstance(text, str):
+            text = [text]
+
+        outputs = []
+
+        for batch in batched(text, n=batch_size):
+            batch_dict = self.tokenizer(batch, max_length=self.max_length, padding=True, truncation=True, return_tensors='pt')
+            output = self.model(**batch_dict)
+            embeddings = self.average_pool(output.last_hidden_state, batch_dict['attention_mask'])
+
+            # embeddings = F.norm(embeddings, p=2, dim=1)
+            # scores = (embeddings[:1] @ embeddings[1:].T) * 100
+
+            embeddings = embeddings.tolist()
+            outputs += embeddings
+        return outputs

src/rag_lanchain.py

@@ -0,0 +1,159 @@
+from typing import TypedDict, Any, List
+import os
+from functools import partial
+
+from langgraph.constants import END
+from qdrant_client import QdrantClient
+
+from langchain_core.prompts import PromptTemplate
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.rate_limiters import InMemoryRateLimiter, BaseRateLimiter
+from langchain_core.tools import tool
+from langchain_core.messages import SystemMessage
+
+from langgraph.graph import START, StateGraph, Graph, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.checkpoint.memory import MemorySaver
+
+from src.database.qdrant_store import QdrantStore
+from src.embeddings import TextEmbedder
+from src.config import EMBEDDING_MODEL, QDRANT_COLLECTION_NAME, CHAT_API_KEY, QDRANT_URL, QDRANT_API_KEY
+
+RAG_PROMPT_STR = """
+You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
+\n
+{context} 
+"""
+RAG_PROMPT = PromptTemplate.from_template(RAG_PROMPT_STR)
+
+embedding_model = TextEmbedder(modelname=EMBEDDING_MODEL)
+
+client = QdrantClient(QDRANT_URL, api_key=QDRANT_API_KEY)
+qdrant_store = QdrantStore(client)
+
+rate_limiter = InMemoryRateLimiter(
+    requests_per_second=0.25,  # <-- Super slow! We can only make a request once every 10 seconds!!
+    check_every_n_seconds=0.1,  # Wake up every 100 ms to check whether allowed to make a request,
+    max_bucket_size=15,  # Controls the maximum burst size.
+)
+
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.0-flash-001",
+    google_api_key=CHAT_API_KEY,
+    temperature=0,
+    max_tokens=None,
+    timeout=None,
+    max_retries=2,
+)
+
+# init_chat_model("google_vertexai:gemini-2.0-flash", rate_limiter=rate_limiter, )
+
+
+class State(TypedDict):
+    question: str
+    context: List[str]
+    answer: str
+
+
+def query_or_respond(state: MessagesState):
+    llm_with_tools = llm.bind_tools([retrieve])
+    response = llm_with_tools.invoke(state["messages"])
+    return {"messages": [response]}
+
+
+@tool
+def retrieve(query: str):
+    """Retrieve information related to a query, specific to the python polars package"""
+    retrieved_docs = []
+    if qdrant_store is not None:
+        query = embedding_model.embed_text(query)
+        retrieved_docs = qdrant_store.get_topk_points_single(query[0], QDRANT_COLLECTION_NAME, k=5)
+    else:
+        retrieved_docs = []
+    return '\n\n'.join(retrieved_docs)
+
+
+def generate(state: MessagesState):
+    recent_tool_messages = []
+    for message in reversed(state["messages"]):
+        if message.type == "tool":
+            recent_tool_messages.append(message)
+        else:
+            break
+    tool_messages = recent_tool_messages[::-1]
+    system_message_content = RAG_PROMPT_STR.format(context=tool_messages)
+    conversation_messages = [
+        message
+        for message in state["messages"]
+        if message.type in ("human", "system")
+           or (message.type == "ai" and not message.tool_calls)
+    ]
+    prompt = [SystemMessage(system_message_content)] + conversation_messages
+
+    response = llm.invoke(prompt)
+    return {"messages": [response]}
+
+
+tools = ToolNode([retrieve])
+
+graph_builder = StateGraph(MessagesState)
+
+graph_builder.add_node(query_or_respond)
+graph_builder.add_node(tools)
+graph_builder.add_node(generate)
+
+graph_builder.set_entry_point("query_or_respond")
+graph_builder.add_conditional_edges(
+    "query_or_respond",
+    tools_condition,
+    {END: END, "tools": "tools"},
+)
+graph_builder.add_edge("tools", "generate")
+graph_builder.add_edge("generate", END)
+
+if __name__ == '__main__':
+    memory = MemorySaver()
+    graph = graph_builder.compile(checkpointer=memory)
+    config = {"configurable": {"thread_id": "def234"}}
+
+    user_input = "Hi there! My name is Will."
+
+    # The config is the **second positional argument** to stream() or invoke()!
+    events = graph.stream(
+        {"messages": [{"role": "user", "content": user_input}]},
+        config,
+        stream_mode="values",
+    )
+    for event in events:
+        event["messages"][-1].pretty_print()
+
+    print(graph.get_state(config))
+    print(memory.get(config))
+
+    user_input = "Remember my name?"
+    config = {"configurable": {"thread_id": "def234"}}
+
+    # The config is the **second positional argument** to stream() or invoke()!
+    events = graph.stream(
+        {"messages": [{"role": "user", "content": user_input}]},
+        config,
+        stream_mode="values",
+    )
+    for event in events:
+        event["messages"][-1].pretty_print()
+
+    print(graph.get_state(config))
+    print(memory.get(config))
+
+    user_input = "Remember my name?"
+    config = {"configurable": {"thread_id": "ddef234"}}
+
+    # The config is the **second positional argument** to stream() or invoke()!
+    events = graph.stream(
+        {"messages": [{"role": "user", "content": user_input}]},
+        config,
+        stream_mode="values",
+    )
+    for event in events:
+        event["messages"][-1].pretty_print()
+

src/retrieval/retrieval.py

@@ -0,0 +1,6 @@
+
+from src.database.qdrant_store import QdrantStore
+from src.embeddings import TextEmbedder
+
+def embed_query(query: str | list[str], ):
+    pass

src/testing.py

@@ -0,0 +1,19 @@
+from bs4 import BeautifulSoup
+from markdown import markdown
+from langchain.text_splitter import MarkdownTextSplitter
+
+
+path = "D:\PycharmProjects\polargs-docu-chat-rag\data\polars-docu\concepts\data-types-and-structures.md"
+
+with open(path, 'r', encoding="utf8") as f_r:
+    test_md = f_r.read()
+
+html = markdown(test_md)
+text = ''.join(BeautifulSoup(html).findAll(text=True))
+
+print(text[:10])
+
+splitter = MarkdownTextSplitter(chunk_size=512, chunk_overlap=64)
+
+docs = splitter.create_documents([text])
+print(docs)

src/utils.py

@@ -0,0 +1,12 @@
+from itertools import islice
+
+
+def batched(iterable, n, *, strict=False):
+    # batched('ABCDEFG', 3) → ABC DEF G
+    if n < 1:
+        raise ValueError('n must be at least one')
+    iterator = iter(iterable)
+    while batch := tuple(islice(iterator, n)):
+        if strict and len(batch) != n:
+            raise ValueError('batched(): incomplete batch')
+        yield batch