initial commit

.gitignore

@@ -1,3 +1,9 @@
+data/
+docs/
+.cursor/
+.chainlit/
+.files/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

chainlit.md

@@ -0,0 +1,14 @@
+# Welcome to Chainlit! 🚀🤖
+
+Hi there, Developer! 👋 We're excited to have you on board. Chainlit is a powerful tool designed to help you prototype, debug and share applications built on top of LLMs.
+
+## Useful Links 🔗
+
+- **Documentation:** Get started with our comprehensive [Chainlit Documentation](https://docs.chainlit.io) 📚
+- **Discord Community:** Join our friendly [Chainlit Discord](https://discord.gg/k73SQ3FyUh) to ask questions, share your projects, and connect with other developers! 💬
+
+We can't wait to see what you create with Chainlit! Happy coding! 💻😊
+
+## Welcome screen
+
+To modify the welcome screen, edit the `chainlit.md` file at the root of your project. If you do not want a welcome screen, just leave this file empty.

chainlit_app.py

@@ -0,0 +1,124 @@
+import chainlit as cl
+from dotenv import load_dotenv
+
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain_core.messages import HumanMessage, AIMessage
+from rag import make_react_agent_graph
+import tiktoken
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.document_loaders import PyMuPDFLoader
+
+# load_dotenv()
+
+
+@cl.on_chat_start
+async def start_chat():
+    settings = { # TODO: These settings might need to be passed to the Langchain model differently
+        "model": "gpt-4o-mini",
+        "temperature": 0.5,
+        "max_tokens": 2000,
+        "frequency_penalty": 0,
+        "presence_penalty": 0,
+    }
+
+    # Initialize Langchain components
+    model = ChatOpenAI(
+        model=settings["model"],
+        temperature=settings["temperature"],
+        # max_tokens=settings["max_tokens"] # ChatOpenAI might not take max_tokens directly here
+    )
+    
+    def tiktoken_len(text):
+        tokens = tiktoken.encoding_for_model("gpt-4o-mini").encode(
+            text,
+        )
+        return len(tokens)
+
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size = 300,
+        chunk_overlap = 0,
+        length_function = tiktoken_len,
+    )
+
+    loader = PyMuPDFLoader("data/PracticalAdviceOnMatrixGames.pdf")
+    docs = loader.load()
+
+    split_chunks = text_splitter.split_documents(docs)
+    embedding_function = OpenAIEmbeddings(model="text-embedding-3-small")
+    # Create a dummy collection. You'll need to populate this with actual documents for RAG to work.
+    vector_store = QdrantVectorStore.from_documents(
+        split_chunks,
+        embedding_function,
+        location=":memory:",
+        collection_name="matrix_game_docs",
+    )
+    # You might want to add some documents here if you have any, e.g.:
+    # vector_store.add_texts(["Some initial context for the agent"])
+
+    # Create the ReAct agent graph
+    # The search_kwargs for the vector store can be customized if needed
+    agent_graph = make_react_agent_graph(model=model, vector_store=vector_store, search_kwargs={"k": 3})
+    
+    cl.user_session.set("agent_graph", agent_graph)
+
+
+@cl.on_message
+async def main(message: cl.Message):
+    agent_graph = cl.user_session.get("agent_graph")
+
+    if not agent_graph:
+        await cl.Message(content="The agent is not initialized. Please restart the chat.").send()
+        return
+
+    conversation_history = cl.chat_context.to_openai()
+
+    msg = cl.Message(content="")
+    # msg.content will be built by streaming tokens.
+    # stream_token will call send() on the first token.
+    async for token, metadata in agent_graph.astream(
+                {'messages': conversation_history},
+                stream_mode="messages"
+            ):
+                
+                if metadata['langgraph_node'] == 'agent':
+                    await msg.stream_token(token.content)
+
+    # try:
+    #     # Use stream_mode="messages" to get LLM tokens as MessageChunk objects
+    #     async for chunk in agent_graph.astream(
+    #         agent_input, {"stream_mode": "messages"}
+    #     ):
+    #         # chunk is expected to be a MessageChunk (e.g., AIMessageChunk)
+    #         if hasattr(chunk, 'content'):
+    #             token = chunk.content
+    #             if token:  # Ensure there's content in the chunk
+    #                 # msg.stream_token will handle sending the message shell on the first call
+    #                 await msg.stream_token(token)
+    #         # else:
+    #             # Handle cases where the chunk might not be a MessageChunk as expected,
+    #             # or if other types of events are streamed in this mode (though less likely for "messages" mode).
+    #             # print(f"Received chunk without content: {chunk}")
+    # except Exception as e:
+    #     print(f"Error during agent stream: {e}")
+    #     await cl.Message(content=f"An error occurred: {str(e)}").send()
+    #     return
+
+    # After the loop, if tokens were streamed, the message content is populated.
+    # A final update might be necessary if other properties of msg need to be set,
+    # or to ensure the stream is properly closed from Chainlit's perspective.
+    if msg.streaming: # msg.streaming is True if stream_token was called
+        await msg.update()
+    elif not msg.content: # If no tokens were streamed and message is still empty
+        # This case might occur if the agent produces no response or an error happened before streaming
+        # (though the try-except should catch errors in astream itself).
+        # Send a default message or handle as an error.
+        # For now, if no content, we don't send an empty message unless explicitly handled.
+        # If msg.send() was never called (because no tokens came), we might need to send something.
+        # However, if there was genuinely no response, sending nothing might be intended.
+        # Let's ensure an empty message isn't sent if no tokens were ever produced.
+        # If an error occurred, it's handled by the except block.
+        # If the agent genuinely returns no content, this will result in no message being sent
+        # beyond the initial empty shell (if it were sent prior to token checks).
+        # Since msg.send() is implicitly called by the first stream_token, if no tokens, no send.
+        pass

rag.py

@@ -0,0 +1,105 @@
+from langgraph.graph import END, StateGraph
+from langgraph.prebuilt import create_react_agent
+from langgraph.checkpoint.memory import InMemorySaver
+
+from langchain_core.vectorstores import VectorStore
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.tools import tool
+from langchain_core.messages import HumanMessage
+from typing import Callable, List, Sequence, Annotated
+from langchain_core.documents import Document
+from typing import Annotated
+from typing_extensions import TypedDict
+from langgraph.graph.message import add_messages
+from langchain_core.documents import Document
+
+class RagState(TypedDict):
+  messages: Annotated[list, add_messages]
+  context: list[Document]
+
+RAG_PROMPT = """\
+Given a provided context and a question, you must answer the question. If you do not know the answer, you must state that you do not know.
+
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:
+"""
+
+rag_prompt_template = ChatPromptTemplate.from_template(RAG_PROMPT)
+
+
+def create_retriever_node(vector_store: VectorStore, search_kwargs: dict = {"k": 5}) -> Callable:
+  def retriever_node(state: RagState) -> RagState:
+    retriever = vector_store.as_retriever(search_kwargs=search_kwargs)
+    retrieved_docs = retriever.invoke(state["messages"][-1].content)
+    return {"context" : retrieved_docs}
+  return retriever_node
+
+
+def create_generator_node(model: BaseChatModel, template: ChatPromptTemplate = rag_prompt_template) -> Callable:
+  generation_chain = template | model
+  def generator_node(state: RagState) -> RagState:
+    response = generation_chain.invoke({"query" : state["messages"][-1].content, "context" : state["context"]})
+    return {"messages" : response}
+  return generator_node
+
+
+def make_rag_graph(model: BaseChatModel, vector_store: VectorStore, template: ChatPromptTemplate = rag_prompt_template, search_kwargs: dict = {"k": 5}) -> StateGraph:
+  retriever_node = create_retriever_node(vector_store, search_kwargs)
+  generator_node = create_generator_node(model, template)
+
+  rag_graph = StateGraph(RagState)
+
+  rag_graph.add_node("retriever", retriever_node)
+  rag_graph.add_node("generator", generator_node)
+
+  rag_graph.set_entry_point("retriever")
+
+  rag_graph.add_edge("retriever", "generator")
+  rag_graph.add_edge("generator", END)
+
+  return rag_graph.compile()
+
+
+# For the ReAct agent, the state is typically managed by the prebuilt agent itself,
+# focusing on the 'messages' list. If a specific state object like RagState is needed
+# for integration, the graph's input/output would need to be adapted.
+# For now, we assume the agent operates on a message-based state.
+
+def create_vector_search_tool(vector_store: VectorStore, search_kwargs: dict) -> Callable:
+  @tool("vector-search")
+  def vector_search_tool(query: str) -> List[str]:
+    """Searches a vector database for the given query and returns relevant document contents."""
+    retriever = vector_store.as_retriever(search_kwargs=search_kwargs)
+    retrieved_docs = retriever.invoke(query)
+    return [doc.page_content for doc in retrieved_docs]
+  return vector_search_tool
+
+
+def make_react_agent_graph(model: BaseChatModel, vector_store: VectorStore, search_kwargs: dict = {"k": 5}):
+  """
+  Creates a StateGraph for a RAG agent that uses the ReAct framework.
+  The agent can formulate its own queries to the vector database.
+  """
+  search_tool = create_vector_search_tool(vector_store, search_kwargs)
+  
+  # A checkpointer is highly recommended for agents to allow them to save state
+  # across multiple steps of thought and action.
+  # checkpointer = InMemorySaver()
+
+  # The prebuilt create_react_agent handles the agent logic and graph compilation.
+  # It uses a default ReAct prompt unless a custom one is provided.
+  # The agent's state revolves around messages.
+  # Input to this agent_graph.invoke should be like: {"messages": [HumanMessage(content="your question")]}
+  agent_graph = create_react_agent(
+      model=model,
+      tools=[search_tool]#,
+      #checkpointer=checkpointer
+  )
+  return agent_graph
+