detectiva / app.py
agentharbor's picture
Update app.py
1cbf972 verified
import os
os.environ['COHERE_API_KEY'] = os.getenv('cohere_ai')
# Create the Cohere chat model
from langchain_cohere.chat_models import ChatCohere
chat = ChatCohere(model="command-r-plus", temperature=0.3)
from langchain_community.tools.tavily_search import TavilySearchResults
os.environ['TAVILY_API_KEY'] = os.getenv('tavily_ai')
internet_search = TavilySearchResults()
internet_search.name = "internet_search"
internet_search.description = "Returns a list of relevant document snippets for a textual query retrieved from the internet."
from langchain_core.pydantic_v1 import BaseModel, Field
class TavilySearchInput(BaseModel):
query: str = Field(description="Query to search the internet with")
internet_search.args_schema = TavilySearchInput
from langchain.agents import Tool
from langchain_experimental.utilities import PythonREPL
python_repl = PythonREPL()
repl_tool = Tool(
name="python_repl",
description="Executes python code and returns the result. The code runs in a static sandbox without interactive mode, so print output or save output to a file.",
func=python_repl.run,
)
repl_tool.name = "python_interpreter"
# from langchain_core.pydantic_v1 import BaseModel, Field
class ToolInput(BaseModel):
code: str = Field(description="Python code to execute.")
repl_tool.args_schema = ToolInput
from langchain.agents import AgentExecutor
from langchain_cohere.react_multi_hop.agent import create_cohere_react_agent
from langchain_core.prompts import ChatPromptTemplate
# Create the prompt
prompt = ChatPromptTemplate.from_template("{input}")
# Create the ReAct agent
agent = create_cohere_react_agent(
llm=chat,
tools=[internet_search, repl_tool],
prompt=prompt,
)
agent_executor = AgentExecutor(agent=agent, tools=[internet_search, repl_tool], verbose=True)
from typing import List, Mapping, Any
from langchain_cohere.common import CohereCitation
def process_data(problem):
output = 'Gemini agent rewriting your query \n\n'
yield output
#rewrite = get_completion(f"Rewrite the user question: {problem} ")
#output += f"Here is your rewritten query: {rewrite} \n\n"
#yield output
output += f"Cohere agent gathering the data from public sources and doing analysis \n\n"
yield output
coh_output = agent_executor.invoke({
"input": f"{problem}"})
print ("Output is",coh_output['output'])
output += f"Final Output: \n\n"+coh_output['output']+"\n\n"
yield output
citations = coh_output['citations']
print (citations)
# Assuming 'citi' is a list of CohereCitation objects
urls = []
for item in citations:
if isinstance(item, CohereCitation) and item.documents:
for doc in item.documents:
if 'url' in doc:
urls.append(doc['url'])
final_urls = list(set(urls))
output += f"Citations: \n\n" + '\n'.join(final_urls)
yield output
return output