done for the night

.gitignore

@@ -1,2 +1,4 @@
 .env
 /__pycache__
+testing.py
+/models

graph.py

@@ -0,0 +1,159 @@
+from typing import Dict, TypedDict, Annotated, Sequence
+from langgraph.graph import Graph, StateGraph, END
+from langgraph.prebuilt import ToolExecutor
+from langchain.schema import StrOutputParser
+from langchain.schema.runnable import RunnablePassthrough
+from langchain_community.tools.tavily_search import TavilySearchResults
+import models
+import prompts
+from helper_functions import format_docs
+from operator import itemgetter
+
+# Define the state structure
+class State(TypedDict):
+    messages: Sequence[str]
+    research_data: Dict[str, str]
+    draft_post: str
+    final_post: str
+
+
+# Research Agent Pieces
+qdrant_research_chain = (
+        {"context": itemgetter("topic") | models.compression_retriever, "topic": itemgetter("topic")}
+        | RunnablePassthrough.assign(context=itemgetter("context"))
+        | {"response": prompts.research_query_prompt  | models.gpt4o_mini | StrOutputParser(), "context": itemgetter("context")}
+    )
+
+# Web Search Agent Pieces
+tavily_tool = TavilySearchResults(max_results=5)
+web_search_chain = (
+    {
+        "topic": itemgetter("topic"),
+        "qdrant_results": itemgetter("qdrant_results"),
+    }
+    | prompts.search_query_prompt
+    | models.gpt4o_mini 
+    | StrOutputParser()
+    | tavily_tool
+    | {
+        "topic": itemgetter("topic"),
+        "qdrant_results": itemgetter("qdrant_results"),
+        "search_results": RunnablePassthrough()
+    }
+    | prompts.summarize_prompt
+    | models.gpt4o_mini 
+    | StrOutputParser()
+)
+
+def query_qdrant(state: State) -> State:
+    # Extract the last message as the input
+    input_text = state["messages"][-1]
+
+    # Run the chain
+    result = qdrant_research_chain.invoke({"topic": input_text})
+
+    # Update the state with the research results
+    state["research_data"]["qdrant_results"] = result
+
+    return state
+
+def web_search(state: State) -> State:
+    # Extract the last message as the topic
+    topic = state["messages"][-1]
+    
+    # Get the Qdrant results from the state
+    qdrant_results = state["research_data"].get("qdrant_results", "No previous results available.")
+    
+    # Run the web search chain
+    result = web_search_chain.invoke({
+        "topic": topic,
+        "qdrant_results": qdrant_results
+    })
+    
+    # Update the state with the web search results
+    state["research_data"]["web_search_results"] = result
+    
+    return state
+
+def research_supervisor(state):
+    # Implement research supervision logic
+    return state
+
+def post_creation(state):
+    # Implement post creation logic
+    return state
+
+def copy_editing(state):
+    # Implement copy editing logic
+    return state
+
+def voice_editing(state):
+    # Implement voice editing logic
+    return state
+
+def post_review(state):
+    # Implement post review logic
+    return state
+
+def writing_supervisor(state):
+    # Implement writing supervision logic
+    return state
+
+def overall_supervisor(state):
+    # Implement overall supervision logic
+    return state
+
+# Create the research team graph
+research_graph = StateGraph(State)
+
+research_graph.add_node("query_qdrant", query_qdrant)
+research_graph.add_node("web_search", web_search)
+research_graph.add_node("research_supervisor", research_supervisor)
+
+research_graph.add_edge("query_qdrant", "research_supervisor")
+research_graph.add_edge("web_search", "research_supervisor")
+research_graph.add_edge("research_supervisor", "query_qdrant")
+research_graph.add_edge("research_supervisor", "web_search")
+research_graph.add_edge("research_supervisor", END)
+
+research_graph.set_entry_point("research_supervisor")
+
+# Create the writing team graph
+writing_graph = StateGraph(State)
+
+writing_graph.add_node("post_creation", post_creation)
+writing_graph.add_node("copy_editing", copy_editing)
+writing_graph.add_node("voice_editing", voice_editing)
+writing_graph.add_node("post_review", post_review)
+writing_graph.add_node("writing_supervisor", writing_supervisor)
+
+writing_graph.add_edge("writing_supervisor", "post_creation")
+writing_graph.add_edge("post_creation", "copy_editing")
+writing_graph.add_edge("copy_editing", "voice_editing")
+writing_graph.add_edge("voice_editing", "post_review")
+writing_graph.add_edge("post_review", "writing_supervisor")
+writing_graph.add_edge("writing_supervisor", END)
+
+writing_graph.set_entry_point("writing_supervisor")
+
+# Create the overall graph
+overall_graph = StateGraph(State)
+
+# Add the research and writing team graphs as nodes
+overall_graph.add_node("research_team", research_graph)
+overall_graph.add_node("writing_team", writing_graph)
+
+# Add the overall supervisor node
+overall_graph.add_node("overall_supervisor", overall_supervisor)
+
+overall_graph.set_entry_point("overall_supervisor")
+
+# Connect the nodes
+overall_graph.add_edge("overall_supervisor", "research_team")
+overall_graph.add_edge("research_team", "overall_supervisor")
+overall_graph.add_edge("overall_supervisor", "writing_team")
+overall_graph.add_edge("writing_team", "overall_supervisor")
+overall_graph.add_edge("overall_supervisor", END)
+
+# Compile the graph
+app = overall_graph.compile()

helper_functions.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Dict, List
 from langchain.agents import AgentExecutor, create_openai_functions_agent
 from langchain_community.document_loaders import PyMuPDFLoader, TextLoader, UnstructuredURLLoader, WebBaseLoader
 from langchain_community.vectorstores import Qdrant
@@ -141,4 +141,7 @@ def create_agent(
     )
     agent = create_openai_functions_agent(llm, tools, prompt)
     executor = AgentExecutor(agent=agent, tools=tools)
-    return executor
+    return executor
+
+def format_docs(docs: List[Dict]) -> str:
+    return "\n\n".join(f"Content: {doc.page_content}\nSource: {doc.metadata.get('source', 'Unknown')}" for doc in docs)

prompts.py

@@ -21,4 +21,45 @@ chat_prompt = ChatPromptTemplate.from_messages([("system", rag_system_prompt_tem
 
 style_guide_path = "./public/CoExperiences Writing Style Guide V1 (2024).pdf"
 style_guide_docs = PyMuPDFLoader(style_guide_path).load()
-style_guide_text = "\n".join([doc.page_content for doc in style_guide_docs])
+style_guide_text = "\n".join([doc.page_content for doc in style_guide_docs])
+
+research_query_prompt = ChatPromptTemplate.from_template("""
+Given a provided context and a topic, compile facts, statistics, quotes, or other related pieces of information that relate to the topic. Make sure to include the source of any such pieces of information in your response.
+
+Context:
+{context}
+
+Topic:
+{topic}
+
+Answer:
+"""
+)
+
+search_query_prompt = ChatPromptTemplate.from_template(
+    """Given the following topic and information from our database, create a search query to find supplementary information:
+
+    Topic: {topic}
+    
+    Information from our database:
+    {qdrant_results}
+    
+    Generate a search query to find additional, up-to-date information that complements what we already know:
+    """
+)
+
+# Create a prompt for summarizing the search results
+summarize_prompt = ChatPromptTemplate.from_template(
+    """Summarize the following search results, focusing on information that is complementary to what we already know from our database. Include sources for each piece of information:
+
+    Topic: {topic}
+    
+    Information from our database:
+    {qdrant_results}
+    
+    Search results:
+    {search_results}
+    
+    Complementary summary with sources:
+    """
+)