Tidy: includes

app.py

@@ -1,21 +1,22 @@
 import asyncio
-from typing import List, Dict, Any
-import chainlit as cl
 import json
 import os
+from dataclasses import dataclass
+from typing import Any, Dict, List
 
+import chainlit as cl
 from dotenv import load_dotenv
-from langchain_openai import ChatOpenAI
-from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import Runnable
+from langchain_openai import ChatOpenAI
+from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain_qdrant import QdrantVectorStore
-from pstuts_rag.loader import load_json_files, load_single_json
 from qdrant_client import QdrantClient
-from dataclasses import dataclass
 
-import pstuts_rag.rag, pstuts_rag.datastore
+import pstuts_rag.datastore
+import pstuts_rag.rag
+from pstuts_rag.loader import load_json_files
 
 
 @dataclass

pstuts_rag/pstuts_rag/agents.py

@@ -0,0 +1,87 @@
+from typing import Any, Callable, List, Optional, TypedDict, Union
+
+from langchain.agents import AgentExecutor, create_openai_functions_agent
+from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+from langchain_core.runnables import Runnable
+from langchain_core.tools import BaseTool
+from langchain_openai import ChatOpenAI
+from langchain_core.language_models import BaseChatModel
+
+from langgraph.graph import END, StateGraph
+
+
+def agent_node(state, agent, name):
+    result = agent.invoke(state)
+    return {"messages": [HumanMessage(content=result["output"], name=name)]}
+
+
+def create_agent(
+    llm: BaseChatModel,
+    tools: list,
+    system_prompt: str,
+):
+    """Create a function-calling agent and add it to the graph."""
+    system_prompt += """
+Work autonomously according to your specialty, using the tools available to you.
+Do not ask for clarification.
+Your other team members (and other teams) will collaborate with you with their own specialties.
+Your first choice should be to use your vector store RAG as the primary source of information. 
+If that does not provide enough context, then use the ArXiv search engine as the first attempt,
+and if that does not provide enough context, then use the Tavily search engine as the second attempt.
+
+You are chosen for a reason! You are one of the following team members: {team_members}.
+"""
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                system_prompt,
+            ),
+            MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+    )
+    agent = create_openai_functions_agent(llm, tools, prompt)
+    executor = AgentExecutor(agent=agent, tools=tools)
+    return executor
+
+
+def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> str:
+    """An LLM-based router."""
+    options = ["FINISH"] + members
+    function_def = {
+        "name": "route",
+        "description": "Select the next role.",
+        "parameters": {
+            "title": "routeSchema",
+            "type": "object",
+            "properties": {
+                "next": {
+                    "title": "Next",
+                    "anyOf": [
+                        {"enum": options},
+                    ],
+                },
+            },
+            "required": ["next"],
+        },
+    }
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder(variable_name="messages"),
+            (
+                "system",
+                "Given the conversation above, who should act next? Or should we FINISH?"
+                "If the last answer was 'I don't know', do not FINISH."
+                " Select one of: {options}",
+            ),
+        ]
+    ).partial(options=str(options), team_members=", ".join(members))
+    return (
+        prompt
+        | llm.bind_functions(functions=[function_def], function_call="route")
+        | JsonOutputFunctionsParser()
+    )

pstuts_rag/pstuts_rag/rag.py

@@ -6,32 +6,25 @@ This module provides the core RAG functionality, including:
 """
 
 import json
-from multiprocessing import Value
 import re
-import uuid
+
 from operator import itemgetter
-from typing import Dict, List, Any, Tuple
+from typing import Any, Dict, List, Tuple
 
+from langchain.prompts import ChatPromptTemplate
 from langchain_core.documents import Document
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage
 from langchain_core.runnables import (
     Runnable,
     RunnableLambda,
     RunnablePassthrough,
 )
-
-from langchain_openai.embeddings import OpenAIEmbeddings
-from langchain_qdrant import QdrantVectorStore
-from qdrant_client import QdrantClient
-
-from langchain.prompts import ChatPromptTemplate
 from langchain_core.vectorstores import VectorStoreRetriever
 from langchain_openai import ChatOpenAI
 
 from .prompt_templates import RAG_PROMPT_TEMPLATES
 
-from langchain_core.language_models.base import BaseLanguageModel
-from langchain_core.messages import AIMessage
-
 
 class RAGChainFactory:
     """Factory class for creating RAG (Retrieval Augmented Generation) chains.
@@ -150,9 +143,7 @@ class RAGChainFactory:
 
     def get_rag_chain(
         self,
-        llm: BaseLanguageModel = ChatOpenAI(
-            model="gpt-4.1-mini", temperature=0
-        ),
+        llm: BaseChatModel = ChatOpenAI(model="gpt-4.1-mini", temperature=0),
     ) -> Runnable:
         """Build and return the complete RAG chain.