agents-in-python/memory-agent/agent.py at main · secrethash/agents-in-python · GitHub

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
import ollama
from tools import calculate, remember, recall

TOOL_REGISTRY = {
    "calculate": calculate,
    "remember": remember,
    "recall": recall,
}

TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Evaluate a mathematical expression.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "Math expression e.g. '12 * 4'",
                    },
                },
                "required": ["expression"],
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "remember",
            "description": "Save important information about the user for future reference.",
            "parameters": {
                "type": "object",
                "properties": {
                    "information": {
                        "type": "string",
                        "description": "The fact or info to remember",
                    }
                },
                "required": ["information"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "recall",
            "description": "Search memory for information relevant to the user's question.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "What to search for in memory"
                    }
                },
                "required": ["query"]
            }
        }
    }
]

conversation_history = []


def run_agent(user_message: str):
    print(f"\n👤 User: {user_message}")

    conversation_history.append({
        "role": "user",
        "content": user_message
    })

    while True:
        response = ollama.chat(
            model="qwen2.5:7b",
            messages=conversation_history,
            tools=TOOL_SCHEMAS
        )

        msg = response.message

        if msg.tool_calls:
            conversation_history.append({
                "role": "assistant",
                "content": msg.content or "",
                "tool_calls": msg.tool_calls
            })

            for tool_call in msg.tool_calls:
                name = tool_call.function.name
                args = tool_call.function.arguments

                print(f"\n🔧 Tool called: {name}({args})")

                result = TOOL_REGISTRY[name](**args) if name in TOOL_REGISTRY else f"Error: tool '{name}' not found."

                print(f"📤 Tool result: {result}")

                conversation_history.append({
                    "role": "tool",
                    "content": str(result)
                })
        else:
            print(f"\n🤖 Agent: {msg.content}")
            conversation_history.append({
                "role": "assistant",
                "content": msg.content
            })
            break