Full graph works. Now frontend and finetuning

pstuts_rag/pstuts_rag/datastore.py

@@ -231,7 +231,7 @@ class DatastoreManager:
             VectorStoreRetriever: The configured retriever
         """
         return self.vector_store.as_retriever(
-            search_kwargs={"k": n_context_docs}
+            search_kwargs={"k": int(n_context_docs)}
         )
 
     def is_ready(self) -> bool:

pstuts_rag/pstuts_rag/graph.py

@@ -118,7 +118,7 @@ def create_agent(
 
 
 def create_tavily_node(
-     name: str = "AdobeHelp", config: Configuration = Configuration() ) -> Callable
+     name: str = "AdobeHelp", config: Configuration = Configuration() ) -> Callable:
     """Initialize tool, agent, and node for Tavily search of helpx.adobe.com.
 
     This function sets up a search agent that can query Adobe Photoshop help topics

pstuts_rag/pstuts_rag/nodes.py

@@ -1,21 +1,29 @@
 # nodes.py
+from enum import Enum
+from typing import Annotated, Any, Callable, Dict, Literal
 
+import asyncio
+import logging
+import operator
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, MessagesState, START, END
 from langgraph.types import Command
-
+from langchain_core.documents import Document
 from langchain_core.runnables import RunnableConfig
-from langchain_core.messages import AnyMessage, HumanMessage
+from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.checkpoint.memory import InMemorySaver
-from pstuts_rag.prompts import NODE_PROMPTS
-from pydantic import BaseModel, Field
+from numpy import add
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_tavily import TavilyExtract
+
+from pydantic import BaseModel, Field, HttpUrl
 
 
 from pstuts_rag.utils import ChatAPISelector
 from pstuts_rag.configuration import Configuration
-
-from enum import Enum
-from typing import Any, Callable, Dict, Literal
+from pstuts_rag.datastore import DatastoreManager
+from pstuts_rag.prompts import NODE_PROMPTS
+from pstuts_rag.rag_for_transcripts import create_transcript_rag_chain
 
 
 class TutorialState(MessagesState):
@@ -23,37 +31,97 @@ class TutorialState(MessagesState):
 
     # next: str
     query: str
-    video_references: set[Any]
-    url_references: set[Any]
+    video_references: Annotated[list[Document], operator.add]
+    url_references: Annotated[list[Dict], operator.add]
+    loop_count: int
+
+
+datastore = DatastoreManager()
+datastore.add_completion_callback(lambda: logging.warning("Loading complete."))
 
 
 def research(state: TutorialState, config: RunnableConfig):
 
-    # retrieve the LLM
-    # configurable = Configuration.from_runnable_config(config)
-    # cls = ChatAPISelector.get(configurable.llm_api, ChatOpenAI)
-    # llm = cls(model=configurable.llm_tool_model)
+    configurable = Configuration.from_runnable_config(config)
+    cls = ChatAPISelector.get(configurable.llm_api, ChatOpenAI)
+    llm = cls(model=configurable.llm_tool_model, temperature=0)
 
-    # # format the prompt
-    # prompt = NODE_PROMPTS["research"]
+    history = [
+        msg.content
+        for msg in state["messages"]
+        if getattr(msg, "role", "") == "ai"
+    ]
 
-    # history = [
-    #     msg.content
-    #     for msg in state["messages"]
-    #     if getattr(msg, "role", "") == "ai"
-    # ]
+    prompt = NODE_PROMPTS["research"].format(
+        history=history, query=state["query"]
+    )
 
-    # prompt = prompt.format(history=history)
+    search_query = llm.invoke([HumanMessage(content=prompt)])
 
-    pass
+    return {
+        "messages": [search_query],
+        "loop_count": state.get("loop_count", 0) + 1,
+    }
 
 
-def search_help(state: TutorialState, config: RunnableConfig):
-    pass
+async def search_help(
+    state: TutorialState, config: RunnableConfig | None = None
+):
+
+    configurable = (
+        Configuration()
+        if not config
+        else Configuration.from_runnable_config(config)
+    )
 
+    cls = ChatAPISelector.get(configurable.llm_api, ChatOpenAI)
+    llm = cls(model=configurable.llm_tool_model, temperature=0)
+    prompt = NODE_PROMPTS["search_summary"]
+
+    adobe_help_search = TavilySearchResults(
+        max_results=2,
+        include_domains=["helpx.adobe.com"],
+        include_answer=True,
+        include_raw_content=True,
+        include_images=True,
+        response_format="content_and_artifact",  # Always returns artifacts
+    )
+    query = state["messages"][-1].content
+    results = await adobe_help_search.ainvoke(query)
 
-def search_rag(state: TutorialState, config: RunnableConfig):
-    pass
+    urls = list(r["url"] for r in results)
+    tool = TavilyExtract(
+        extract_depth="basic",
+        include_images=False,
+    )
+
+    results = await tool.ainvoke({"urls": urls})
+
+    if "results" in results:
+        all_text = list(r["raw_content"] for r in results["results"])
+    else:
+        all_text = []
+
+    prompt = prompt.format(
+        query=query,
+        text="\n***\n".join(all_text),
+    )
+
+    url_summary = await llm.ainvoke([HumanMessage(content=prompt)])
+
+    return {"messages": [url_summary], "url_references": results["results"]}
+
+
+async def search_rag(state: TutorialState, config: RunnableConfig):
+
+    chain = create_transcript_rag_chain(datastore, config)
+
+    response = await chain.ainvoke({"question": state["messages"][-1].content})
+
+    return {
+        "messages": [response],
+        "video_references": response.additional_kwargs["context"],
+    }
 
 
 def join(state: TutorialState, config: RunnableConfig):
@@ -71,6 +139,11 @@ class YesNoDecision(BaseModel):
     decision: Literal["yes", "no"] = Field(description="Yes or no decision.")
 
 
+class URLReference(BaseModel):
+    summary: str
+    url: HttpUrl
+
+
 def route_is_relevant(
     state: TutorialState, config: RunnableConfig
 ) -> Command[Literal["research", "write_answer"]]:
@@ -82,25 +155,103 @@ def route_is_relevant(
         YesNoDecision
     )
 
+    human_messages = [
+        msg.content
+        for msg in state["messages"]
+        if isinstance(msg, HumanMessage)
+    ]
+
+    if len(human_messages) > 0:
+        query = human_messages[-1]
+    else:
+        query = state["query"]
+
     # format the prompt
-    prompt = NODE_PROMPTS["relevance"].format(query=state["query"])
+    prompt = NODE_PROMPTS["relevance"].format(query=query)
 
     relevance = llm.invoke([HumanMessage(content=prompt)])
     where = "research" if relevance.decision == "yes" else "write_answer"
-    answer = f"Query is {'not' if relevance.decision == 'no' else ''} relevant to Photoshop."
+    answer = (
+        f"Query is {'not' if relevance.decision == 'no' else ''} "
+        "relevant to Photoshop."
+    )
     return Command(
-        update={"messages": {"role": "ai", "content": answer}},
+        update={"messages": [AIMessage(content=answer)], "query": query},
         goto=where,
     )
 
 
+class IsComplete(BaseModel):
+    decision: Literal["yes", "no"] = Field(description="Yes or no decision.")
+    new_query: str = Field(description="Query for additional research.")
+
+
 def route_is_complete(
     state: TutorialState, config: RunnableConfig
-) -> Literal["yes", "no"]:
-    if True:
-        return "yes"
-    else:
-        return "no"
+) -> Command[Literal["research", "write_answer"]]:
+
+    # retrieve the LLM
+    configurable = Configuration.from_runnable_config(config)
+
+    if state["loop_count"] >= int(configurable.max_research_loops):
+        return Command(
+            update={
+                "messages": [
+                    AIMessage(
+                        content="Research loop count is too large. Do your best with what you have."
+                    )
+                ]
+            },
+            goto="write_answer",
+        )
+
+    cls = ChatAPISelector.get(configurable.llm_api, ChatOpenAI)
+    llm = cls(model=configurable.llm_tool_model).with_structured_output(
+        YesNoDecision
+    )
+
+    ai_messages = list(
+        msg.content for msg in state["messages"] if isinstance(msg, AIMessage)
+    )
+
+    # format the prompt
+    prompt = NODE_PROMPTS["completeness"].format(
+        query=state["query"], responses="\n\n".join(ai_messages)
+    )
+
+    completeness = llm.invoke([HumanMessage(content=prompt)])
+    where = "write_answer" if "yes" in completeness.decision else "research"
+
+    # Convert YesNoDecision to AIMessage
+    decision_message = AIMessage(
+        content=f"Research completeness: {completeness.decision}"
+    )
+
+    return Command(
+        update={"messages": [decision_message]},
+        goto=where,
+    )
+
+
+def write_answer(state: TutorialState, config: RunnableConfig):
+
+    # retrieve the LLM
+    configurable = Configuration.from_runnable_config(config)
+    cls = ChatAPISelector.get(configurable.llm_api, ChatOpenAI)
+    llm = cls(model=configurable.llm_tool_model)
+
+    ai_messages = list(
+        msg.content for msg in state["messages"] if isinstance(msg, AIMessage)
+    )
+
+    # format the prompt
+    prompt = NODE_PROMPTS["final_answer"].format(
+        query=state["query"], responses="\n\n".join(ai_messages)
+    )
+
+    final_answer = llm.invoke([HumanMessage(content=prompt)])
+
+    return {"messages": [final_answer]}
 
 
 graph_builder = StateGraph(TutorialState)
@@ -119,17 +270,16 @@ graph_builder.add_node(write_answer)
 #     {"yes": research.__name__, "no": write_answer.__name__},
 # )
 graph_builder.add_node(route_is_relevant)
+graph_builder.add_node(route_is_complete, defer=True)
+
 graph_builder.add_edge(START, route_is_relevant.__name__)
 graph_builder.add_edge(research.__name__, search_help.__name__)
 graph_builder.add_edge(research.__name__, search_rag.__name__)
-graph_builder.add_edge(search_help.__name__, join.__name__)
-graph_builder.add_edge(search_rag.__name__, join.__name__)
-graph_builder.add_conditional_edges(
-    join.__name__,
-    route_is_complete,
-    {"no": research.__name__, "yes": write_answer.__name__},
-)
+graph_builder.add_edge(search_help.__name__, route_is_complete.__name__)
+graph_builder.add_edge(search_rag.__name__, route_is_complete.__name__)
+
 graph_builder.add_edge(write_answer.__name__, END)
 
 
 graph = graph_builder.compile()
+asyncio.run(datastore.from_json_globs(Configuration().transcript_glob))

pstuts_rag/pstuts_rag/prompts.py

@@ -164,3 +164,67 @@ is relevant to Adobe Photoshop, otherwise no.
 
 Relevant?
 """
+
+NODE_PROMPTS[
+    "search_summary"
+] = """
+<QUERY>
+{query}
+</QUERY>
+<WEBSITE_TEXT>
+{text}
+</WEBSITE_TEXT>
+
+<TASK>
+Use WEBSITE_TEXT to produce a summarized
+answer to the QUERY.
+
+Aim for the audience at a level of an advanced high school student.
+Do not invent material that is not in the text.
+
+Your output should be at most 200 words long.
+</TASK>
+"""
+
+NODE_PROMPTS[
+    "completeness"
+] = """
+<QUERY>
+{query}
+</QUERY>
+<RESEARCH>
+{responses}
+</RESEARCH>
+
+<TASK>
+Your goal is to evaluate if RESEARCH is sufficiently detailed to provide a comprehensive
+and clear answer for QUERY.
+
+If the RESEARCH is sufficiently complete, state "yes" as your decision. 
+
+If new terms were introduced in RESEARCH that are not sufficiently explained,
+or the QUERY is not sufficiently addressed, response as "no".
+</TASK>
+
+<FINAL_CHECK>
+Your response must be either "yes" or "no".
+</FINAL_CHECK>
+"""
+
+NODE_PROMPTS[
+    "final_answer"
+] = """
+<QUERY>
+{query}
+</QUERY>
+<RESEARCH>
+{responses}
+</RESEARCH>
+
+<TASK>
+Use the content in RESEARCH to provide a detailed answer to the QUERY.
+Do not add the material, fully ground yourself in the research context.
+
+End your response with "I hope you're happy!".
+</TASK>
+"""

pstuts_rag/pstuts_rag/rag_for_transcripts.py

@@ -58,10 +58,10 @@ def post_process_response(
         else answer.content
     )
     # Only append references if the model provided a substantive answer
-    if "I don't know" not in answer.content:
-        text_w_references = "\n".join(
-            [str(text_w_references), "**REFERENCES**", references]
-        )
+    # if "I don't know" not in answer.content:
+    #     text_w_references = "\n".join(
+    #         [str(text_w_references), "**REFERENCES**", references]
+    #     )
 
     # Create new message with references and preserve original context metadata
     output: AIMessage = answer.model_copy(

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "langchain-ollama>=0.3.2",
     "simsimd>=6.2.1",
     "langgraph-cli[inmem]>=0.1.55",
+    "langchain-tavily>=0.2.0",
 ]
 authors = [{ name = "Marko Budisic", email = "[email protected]" }]
 license = "MIT"

temp_function.txt

@@ -0,0 +1,6 @@
+def enter_chain(message: str):
+    results = {
+        "messages": [HumanMessage(content=message)],
+        "team_members": ["VideoArchiveSearch", "AdobeHelp"],
+    }
+    return results

uv.lock

@@ -1879,6 +1879,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/68/01/22dad84373ba282237a3351547443c9c94c39fe75f71a1759f97cfa89725/langchain_qdrant-0.2.0-py3-none-any.whl", hash = "sha256:8eab5b8a553204ddb809d8183a6f1bc12fc265688592d9d897388f6939c79bf8", size = 23406 },
 ]
 
+[[package]]
+name = "langchain-tavily"
+version = "0.2.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "aiohttp" },
+    { name = "langchain" },
+    { name = "langchain-core" },
+    { name = "mypy" },
+    { name = "requests" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/df/63/e7c41f837914806b3c255c4c46d0948528101279656a523b7e11be740e06/langchain_tavily-0.2.0.tar.gz", hash = "sha256:b400525d6d2c28902d2acb25af28751aa1a9a1f99c7880eea4d701f3993736fb", size = 19813 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b5/a7/2e59086df6006ac09a8d8d8f43683ff2f84608d69984bf1593c92faeefb0/langchain_tavily-0.2.0-py3-none-any.whl", hash = "sha256:a5b780f96c80d5a3e7c933da2d603cb26ba94b10f7c1ac4b89ce5b123c7541b4", size = 23580 },
+]
+
 [[package]]
 name = "langchain-text-splitters"
 version = "0.3.8"
@@ -3741,6 +3757,7 @@ dependencies = [
     { name = "langchain-ollama" },
     { name = "langchain-openai" },
     { name = "langchain-qdrant" },
+    { name = "langchain-tavily" },
     { name = "langgraph" },
     { name = "langgraph-cli", extra = ["inmem"] },
     { name = "langsmith" },
@@ -3810,6 +3827,7 @@ requires-dist = [
     { name = "langchain-ollama", specifier = ">=0.3.2" },
     { name = "langchain-openai" },
     { name = "langchain-qdrant", specifier = ">=0.2.0" },
+    { name = "langchain-tavily", specifier = ">=0.2.0" },
     { name = "langgraph", specifier = ">=0.2.55" },
     { name = "langgraph-cli", extras = ["inmem"], specifier = ">=0.1.55" },
     { name = "langsmith", specifier = ">=0.0.50" },