File size: 3,826 Bytes
0fbe2c5
 
 
 
 
 
 
 
 
 
 
 
 
 
b904ff2
0fbe2c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b904ff2
 
 
 
0fbe2c5
 
b904ff2
0fbe2c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import base64
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_community.tools import DuckDuckGoSearchResults
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
import wikipediaapi
import json
import asyncio
import aiohttp
from langchain_core.tools import tool
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_tavily import TavilySearch

import requests

system_prompt = """You are a helpful assistant tasked with answering questions using a set of tools.
Now, I will ask you a question. Report your thoughts, and finish your answer with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER].
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
"""

api_key = os.getenv("OPENAI_API_KEY")

model = ChatOpenAI(model="gpt-4o-mini", api_key=api_key, temperature=0)

@tool
def search_wiki(query: str, max_results: int = 3) -> str:
    """
    Searches Wikipedia for the given query and returns a maximum of 'max_results'
    relevant article summaries, titles, and URLs.

    Args:
        query (str): The search query for Wikipedia.
        max_results (int): The maximum number of search results to retrieve (default is 3).

    Returns:
        str: A JSON string containing a list of dictionaries, where each dictionary
             represents a Wikipedia article with its title, summary, and URL.
             Returns an empty list if no results are found or an error occurs.
    """

    language_code = 'en'

    headers={'User-Agent': 'LangGraphAgent/1.0 ([email protected])'}

    base_url = 'https://api.wikimedia.org/core/v1/wikipedia/'
    endpoint = '/search/page'
    url = base_url + language_code + endpoint
    parameters = {'q': query, 'limit': max_results}
    response = requests.get(url, headers=headers, params=parameters)
    response = json.loads(response.text)
    return json.dumps(response, indent=2)


tavily_search_tool = TavilySearch(
    max_results=5,
    topic="general",
)

tools = [
    tavily_search_tool,
    search_wiki
]

def build_graph():
    """Build the graph"""
    # Bind tools to LLM
    llm_with_tools = model.bind_tools(tools)

    # Node
    def assistant(state: MessagesState):
        """Assistant node"""
        return {"messages": [llm_with_tools.invoke(state["messages"])]}
    
    builder = StateGraph(MessagesState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
    )
    builder.add_edge("tools", "assistant")

    # Compile graph
    return builder.compile()


# test agent
if __name__ == "__main__":
    question = "When was St. Thomas Aquinas born?"
    # Build the graph
    graph = build_graph()
    # Run the graph
    messages = [
        SystemMessage(
            content=system_prompt
        ),
        HumanMessage(
            content=question
        )]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()