File size: 4,600 Bytes
ebec9e2
 
 
 
 
 
 
 
 
 
 
 
 
 
b15267e
 
 
 
ebec9e2
 
 
 
 
 
 
 
 
 
 
 
 
b15267e
 
 
 
ebec9e2
 
43b41d6
b15267e
 
ebec9e2
 
b15267e
ebec9e2
b15267e
 
ebec9e2
b15267e
ebec9e2
b15267e
ebec9e2
b15267e
ebec9e2
 
 
 
b15267e
ebec9e2
b15267e
ebec9e2
b15267e
ebec9e2
 
b15267e
ebec9e2
73caeb8
b15267e
ebec9e2
 
 
 
 
 
 
 
 
 
 
 
 
b15267e
ebec9e2
 
 
 
b15267e
 
ebec9e2
 
 
 
 
 
 
 
 
 
 
 
b15267e
ebec9e2
b15267e
ebec9e2
 
b15267e
ebec9e2
73caeb8
b15267e
ebec9e2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from typing import Dict, Any, List
from langgraph.graph import Graph, StateGraph
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field
from duckduckgo_search import DDGS

class CalculatorInput(BaseModel):
    operation: str = Field(..., description="The operation to perform (add, subtract, multiply, divide)")
    numbers: List[float] = Field(..., description="List of numbers to perform the operation on")

class CalculatorOutput(BaseModel):
    result: float = Field(..., description="The result of the calculation")
    operation: str = Field(..., description="The operation that was performed")

class CalculatorState(BaseModel):
    input: CalculatorInput
    output: CalculatorOutput = None

class SearchInput(BaseModel):
    query: str = Field(..., description="The search query to look up")
    max_results: int = Field(default=3, description="Maximum number of results to return")

class SearchResult(BaseModel):
    title: str = Field(..., description="Title of the search result")
    link: str = Field(..., description="URL of the search result")
    snippet: str = Field(..., description="Brief description of the search result")

class SearchOutput(BaseModel):
    results: List[SearchResult] = Field(..., description="List of search results")
    query: str = Field(..., description="The original search query")

class SearchState(BaseModel):
    input: SearchInput
    output: SearchOutput = None

def create_calculator_tool() -> Graph:
    """Creates a calculator tool using LangGraph that can perform basic arithmetic operations."""
    print("Creating calculator tool")
    def calculator_function(state: CalculatorState) -> CalculatorState:
        if len(state.input.numbers) < 2:
            raise ValueError("At least two numbers are required for calculation")
        
        result = state.input.numbers[0]
        
        for num in state.input.numbers[1:]:
            if state.input.operation == "add":
                result += num
            elif state.input.operation == "subtract":
                result -= num
            elif state.input.operation == "multiply":
                result *= num
            elif state.input.operation == "divide":
                if num == 0:
                    raise ValueError("Cannot divide by zero")
                result /= num
            else:
                raise ValueError(f"Unsupported operation: {state.input.operation}")
        
        state.output = CalculatorOutput(
            result=result,
            operation=state.input.operation
        )
        
        return state

    # Create the graph with state schema
    workflow = StateGraph(state_schema=CalculatorState)
    
    # Add the calculator tool node
    workflow.add_node("calculator", ToolNode(calculator_function))
    
    # Set the entry and exit points
    workflow.set_entry_point("calculator")
    workflow.set_finish_point("calculator")
    
    return workflow.compile()

def create_search_tool() -> Graph:
    """Creates a search tool using DuckDuckGo that can search for information online."""
    
    def search_function(state: SearchState) -> SearchState:
        # Initialize DuckDuckGo search
        with DDGS() as ddgs:
            # Perform the search
            search_results = list(ddgs.text(
                state.input.query,
                max_results=state.input.max_results
            ))
            
            # Convert results to our model
            results = [
                SearchResult(
                    title=result["title"],
                    link=result["link"],
                    snippet=result["body"]
                )
                for result in search_results
            ]
            
            state.output = SearchOutput(
                results=results,
                query=state.input.query
            )
            
            return state

    # Create the graph with state schema
    workflow = StateGraph(state_schema=SearchState)
    
    # Add the search tool node
    workflow.add_node("search", ToolNode(search_function))
    
    # Set the entry and exit points
    workflow.set_entry_point("search")
    workflow.set_finish_point("search")
    
    return workflow.compile()

# Example usage:
# if __name__ == "__main__":
#     # Create the calculator tool
#     calculator = create_calculator_tool()
    
#     # Example calculation
#     result = calculator.invoke({
#         "input": {
#             "operation": "add",
#             "numbers": [1, 2, 3, 4]
#         }
#     })
    
#     print(f"Result: {result['output'].result}")