File size: 4,554 Bytes
91066f0 176a5a5 e95479f 91066f0 c23f1e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import streamlit as st
import os
from dotenv import load_dotenv
from tavily import TavilyClient
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import AgentExecutor, Tool
from langchain.agents import create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.graph import END, Graph
load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
llm = ChatGoogleGenerativeAI(
model="gemini-2.0-flash",
temperature=0.7,
google_api_key=GOOGLE_API_KEY
)
tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
def tavily_search(query: str):
try:
result = tavily_client.search(
query=query,
search_depth="advanced",
include_raw_content=True
)
return str(result)
except Exception as e:
return f"Error searching Tavily: {str(e)}"
tavily_tool = Tool(
name="tavily_search",
func=tavily_search,
description="A premium search engine that returns high-quality, up-to-date information."
)
def create_researcher_agent():
prompt = ChatPromptTemplate.from_messages([
("system", """You are a senior research analyst. Your job:
1. Use Tavily for authoritative info.
2. Cross-validate across multiple sources.
3. Extract key facts, stats & quotes.
4. Always include source URLs."""),
("user", "{input}"),
MessagesPlaceholder(variable_name="agent_scratchpad")
])
agent = create_tool_calling_agent(llm, [tavily_tool], prompt)
return AgentExecutor(agent=agent, tools=[tavily_tool])
def generate_final_report(research_output: str) -> str:
prompt = ChatPromptTemplate.from_messages([
("system", """You are a technical writer. Structure and present the research:
- Use markdown formatting
- Academic tone
- Include proper citations
- Highlight insights.
- Make sure to answer the user's original query completely."""),
("user", "{input}")
])
formatted_prompt = prompt.format_messages(input=research_output)
return llm.invoke(formatted_prompt).content
def create_workflow():
workflow = Graph()
def research_node(state):
agent = create_researcher_agent()
output = agent.invoke({"input": state["query"], "agent_scratchpad": []})
return {"research_output": output["output"]}
def write_node(state):
report = generate_final_report(state["research_output"])
return {"final_report": report}
workflow.add_node("research", research_node)
workflow.add_node("write", write_node)
workflow.add_edge("research", "write")
workflow.add_edge("write", END)
workflow.set_entry_point("research")
return workflow
def main():
st.set_page_config(page_title="AI Research Assistant", page_icon="🔍", layout="wide")
st.title(" AI Research Assistant")
st.markdown("""Enter a query and let the AI prepare a research-backed report:
- Searches authoritative sources
- Validates information
- Delivers a professional markdown report.
""")
query = st.text_area("Enter your research query:",
placeholder="e.g., 'Impact of AI on healthcare jobs in 2024'",
height=100)
if st.button("Research", type="primary"):
if not query:
st.warning("Please enter a research query first.")
return
if not GOOGLE_API_KEY or not TAVILY_API_KEY:
st.error(" API keys missing! Check your `.env` file.")
return
with st.spinner("Researching your topic..."):
try:
workflow = create_workflow()
app = workflow.compile()
final_output = app.invoke({"query": query})
with st.expander(" Research Report", expanded=True):
if "final_report" in final_output:
st.markdown(final_output["final_report"])
elif "research_output" in final_output:
st.markdown("### Research Findings (Unformatted)")
st.markdown(final_output["research_output"])
else:
st.warning("No results were generated")
st.success(" Research completed!")
except Exception as e:
st.error(f" An error occurred: {str(e)}")
st.exception(e)
if __name__ == "__main__":
main() |