from dotenv import load_dotenv from langchain_openai import ChatOpenAI,OpenAIEmbeddings from pydantic import BaseModel, Field from langchain_core.messages import HumanMessage, AIMessage, SystemMessage import json import uuid from datetime import datetime, timezone, timedelta from langchain_community.tools import DuckDuckGoSearchResults from langchain_community.utilities import DuckDuckGoSearchAPIWrapper from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_chroma import Chroma from typing import Literal, TypedDict from langgraph.graph import StateGraph, START, END from IPython.display import display,Image from langgraph.checkpoint.memory import InMemorySaver from graph_states import DataExtractionState,Queries,Extractor,Analyzer from langgraph.checkpoint.memory import InMemorySaver import asyncio load_dotenv() research_model = openai_model = ChatOpenAI(model="gpt-5-nano",reasoning_effort="minimal") embedding_model = OpenAIEmbeddings( model="text-embedding-3-small", ) def generate_subqueries(state): prompt = f"""You are an assistant that expands a user’s query into a set of focused, well-structured sub-queries for web search and research. Goal: Given any user query, generate a list of 5 sub-queries that explore the topic from multiple perspectives, ensure comprehensive coverage, and surface context the user may not have considered. Instructions: Understand the user’s original query and identify its core intent. Queries must be upto date with current time and date {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')}. Break it into smaller, precise, search-ready sub-queries. Include different angles such as: Background and foundational concepts Technical details Related problems or challenges Historical context Current trends or recent developments Tools, methods, or frameworks Comparisons and alternatives Case studies and practical examples Controversies, limitations, or open research questions Avoid redundancy; each sub-query should explore a unique dimension. Use concise, clear phrasing suitable for search engines.""" messages = [ SystemMessage(content=prompt), HumanMessage(content=state["query"]) ] structured_model = research_model.with_structured_output(Queries) output = structured_model.invoke(messages) return {"subqueries":output.queries,"vector_store":str(uuid.uuid4())} def search_subqueries(state): wrapper = DuckDuckGoSearchAPIWrapper(time="d", max_results=10) search = DuckDuckGoSearchResults(api_wrapper=wrapper, source="news",output_format='list') # search = tool search_results = [] subqueries = state["subqueries"] for query in subqueries: result=search.invoke(query) search_results.append(json.dumps(result)) return {"results":search_results} def create_vectorstore_ddg(state): results = state["results"] for result in results: documents = [] result=json.loads(result) for item in result: page_content = item.get("snippet", "") title = item.get("title", "No Title") link = item.get("link", "") doc = Document( page_content=page_content, metadata={"source": link, "title": title} ) documents.append(doc) splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder( encoding_name="cl100k_base", chunk_size=128, chunk_overlap=24 ) texts = splitter.split_documents(documents) vector_store = Chroma( collection_name=state["vector_store"], embedding_function=embedding_model, persist_directory="./chroma_langchain_db", tenant="default_tenant", ) vector_store.add_documents(texts) return state async def batch_extract(state): prompt = """Role: You are a precision extraction model. Your job is to read the provided web-scraped content and extract only the information relevant to the user’s query. You do not interpret, judge, or analyze correctness—only extract. Instructions: Read the user’s query. Read the provided scraped content. Identify all information that is directly relevant to fulfilling the query. Remove everything irrelevant: filler text, ads, disclaimers, unrelated sections, partial sentences, duplicates, noise, formatting artifacts, navigation menus, etc. When information appears multiple times, keep the clearest version. Preserve original factual wording when possible. If a required piece of information is missing, output "not found" for that field. Do not guess or invent content under any circumstances. Output the result in a compact, structured format. Inputs: User Query: {QUERY} Your task: Return only the structured extraction described above.""" # structured_model = gemini_model.with_structured_output(Extractor.model_json_schema()) structured_model = research_model.with_structured_output(Extractor.model_json_schema()) batches = [] queries = state["subqueries"] results = state["results"] for query,result in zip(queries,results): message = [ SystemMessage(content=prompt.format(QUERY=query,CONTENT=result)), HumanMessage(content=f"Scraped {result}"), ] batches.append(message) results = await structured_model.abatch(batches) return {"subquery_extracted_data":results} async def analyzer(state): structured_model = research_model.with_structured_output(Analyzer) ANALYZER_PROMPT = """ You are an analysis model. Your job is to decide whether the extracted information satisfies the query. Rules: 1. Use ONLY the extracted information. Do not add or guess any facts. 2. Determine if the query is satisfied: True or False. 3. Provide short, logical reasoning. 4. If any required info is missing or uncertain, list it. Otherwise return an empty list. 5. Schema explanation: {{ "query_statisfied": , "reasoning": "", "missing_or_uncertain_information": [ "", ... ] }} User Query: {query} Extracted Information: {extracted} """ if state.get("subquery_extracted_data",None): data = state["subquery_extracted_data"] else: data = state["results"] batches = [] for query,extracted in zip(state["subqueries"],data): message = [ SystemMessage(content=ANALYZER_PROMPT.format(query=query,extracted=json.dumps(extracted))) ] batches.append(message) results = await structured_model.abatch(batches) return {"subquery_reviews": results} async def query_refiner(state): context_dicts = [] for review in state["subquery_reviews"]: context_dicts.append(review.model_dump()) context_json_string = json.dumps(context_dicts) prompt = f""" # Role You are an expert Query Refinement Engine. Your goal is to analyze failed search attempts and generate a single, highly optimized search query to retrieve the missing information. Queries must be upto date with current time and date {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')}. # Input Context You will receive a list of JSON objects representing the analysis of previous search attempts. Each object follows this schema: - `query_statisfied` (bool): Whether the intent was met. - `query` (str): The original query attempted. - `reasoning` (str): Context on why it failed or succeeded. - `missing_or_uncertain_information` (list[str]): Specific data points that are still needed. # Task Iterate through the input list. **Only** process entries where `query_statisfied` is `False`. For each unsatisfied entry, generate **ONE** (1) refined query. 1. **Consolidate:** If multiple pieces of information are missing for a single entry, attempt to combine them into one specific, high-density query. 2. **Precision:** Use the `reasoning` field to identify why the previous attempt failed (e.g., broadness, wrong domain) and adjust keywords accordingly. 3. **Self-Correction:** If the previous query was a question, try a keyword string. If it was keywords, try a natural language question. # Constraints - **Do not** generate multiple queries for a single unsatisfied entry unless the missing information items are completely unrelated topics. - **Maximum limit:** Strictly produce no more than 1 refined query per unsatisfied input entry. - If a query failed due to being "too broad", the refined query must be narrower. # Context: {context_json_string} """ structured_model = research_model.with_structured_output(Queries) response = structured_model.invoke([SystemMessage(content=prompt)]) state["subqueries"] = response.queries state["query_limit"]-=1 return state async def refinement_check(state)->Literal["Query Refiner",END]: for review in state["subquery_reviews"]: if (not review.query_statisfied) and state["query_limit"]>0: return "Query Refiner" return END def check_result_size(state)->Literal["Batch Extract","Analyzer"]: results = state["results"] data = "".join(results) if len(data)>12000: return "Batch Extract" else: return "Analyzer" async def create_research_graph(): blueprint = StateGraph(DataExtractionState) blueprint.add_node("Subqueries node",generate_subqueries) blueprint.add_node("Search Subqueries",search_subqueries) blueprint.add_node("Create VectorStore",create_vectorstore_ddg) blueprint.add_node("Batch Extract",batch_extract) blueprint.add_node("Analyzer",analyzer) blueprint.add_node("Query Refiner",query_refiner) blueprint.add_edge(START,"Subqueries node") blueprint.add_edge("Subqueries node","Search Subqueries") blueprint.add_edge("Search Subqueries","Create VectorStore") blueprint.add_edge("Batch Extract","Analyzer") blueprint.add_conditional_edges("Create VectorStore",check_result_size) blueprint.add_conditional_edges("Analyzer",refinement_check) blueprint.add_edge("Query Refiner","Search Subqueries") graph = blueprint.compile(checkpointer=InMemorySaver()) display(Image(graph.get_graph().draw_mermaid_png())) return graph if __name__ == "__main__": async def main(): research_graph = await create_research_graph() thread_id = str(uuid.uuid4()) config = {"configurable":{"thread_id":thread_id}} initial_state: DataExtractionState = { "query": "What are the latest advancements in renewable energy technologies in 2024?", "subqueries": [], "subquery_extracted_data": None, "subquery_reviews": [], "results": [], "vector_store": "", "query_limit":1 } final_state = await research_graph.ainvoke(initial_state) print("Final State:", final_state) asyncio.run(main())