1 files changed, 94 insertions, 8 deletions

diff --git a/agent/pw_agent.go b/agent/pw_agent.go
index 8c1c2bf..2807331 100644
--- a/agent/pw_agent.go
+++ b/agent/pw_agent.go
@@ -1,5 +1,11 @@
 package agent
 
+import (
+	"encoding/json"
+	"gf-lt/models"
+	"strings"
+)
+
 // PWAgent: is AgenterA type agent (enclosed with tool chaining)
 // sysprompt explain tools and how to plan for execution
 type PWAgent struct {
@@ -7,11 +13,16 @@ type PWAgent struct {
 	sysprompt string
 }
 
-// NewWebAgentB creates a WebAgentB that uses the given formatting function
+// NewPWAgent creates a PWAgent with the given client and system prompt
 func NewPWAgent(client *AgentClient, sysprompt string) *PWAgent {
 	return &PWAgent{AgentClient: client, sysprompt: sysprompt}
 }
 
+// SetTools sets the tools available to the agent
+func (a *PWAgent) SetTools(tools []models.Tool) {
+	a.tools = tools
+}
+
 func (a *PWAgent) ProcessTask(task string) []byte {
 	req, err := a.FormFirstMsg(a.sysprompt, task)
 	if err != nil {
@@ -25,16 +36,91 @@ func (a *PWAgent) ProcessTask(task string) []byte {
 			a.Log().Error("failed to process the request", "error", err)
 			return []byte("failed to process the request; err: " + err.Error())
 		}
-		toolCall, hasToolCall := findToolCall(resp)
+		execTool, toolCallID, hasToolCall := findToolCall(resp)
 		if !hasToolCall {
 			return resp
 		}
-		// check resp for tool calls
-		// make tool call
-		// add tool call resp to body
-		// send new request too lmm
-		tooResp := toolCall(resp)
-		req, err = a.FormMsg(toolResp)
+
+		a.setToolCallOnLastMessage(resp, toolCallID)
+
+		toolResp := string(execTool())
+		req, err = a.FormMsgWithToolCallID(toolResp, toolCallID)
+		if err != nil {
+			a.Log().Error("failed to form next message", "error", err)
+			return []byte("failed to form next message; err: " + err.Error())
+		}
 	}
 	return nil
 }
+
+func (a *PWAgent) setToolCallOnLastMessage(resp []byte, toolCallID string) {
+	if toolCallID == "" {
+		return
+	}
+
+	var genericResp map[string]interface{}
+	if err := json.Unmarshal(resp, &genericResp); err != nil {
+		return
+	}
+
+	var name string
+	var args map[string]string
+
+	if choices, ok := genericResp["choices"].([]interface{}); ok && len(choices) > 0 {
+		if firstChoice, ok := choices[0].(map[string]interface{}); ok {
+			if message, ok := firstChoice["message"].(map[string]interface{}); ok {
+				if toolCalls, ok := message["tool_calls"].([]interface{}); ok && len(toolCalls) > 0 {
+					if tc, ok := toolCalls[0].(map[string]interface{}); ok {
+						if fn, ok := tc["function"].(map[string]interface{}); ok {
+							name, _ = fn["name"].(string)
+							argsStr, _ := fn["arguments"].(string)
+							json.Unmarshal([]byte(argsStr), &args)
+						}
+					}
+				}
+			}
+		}
+	}
+
+	if name == "" {
+		content, _ := genericResp["content"].(string)
+		name = extractToolNameFromText(content)
+	}
+
+	lastIdx := len(a.chatBody.Messages) - 1
+	if lastIdx >= 0 {
+		a.chatBody.Messages[lastIdx].ToolCallID = toolCallID
+		if name != "" {
+			argsJSON, _ := json.Marshal(args)
+			a.chatBody.Messages[lastIdx].ToolCall = &models.ToolCall{
+				ID:   toolCallID,
+				Name: name,
+				Args: string(argsJSON),
+			}
+		}
+	}
+}
+
+func extractToolNameFromText(text string) string {
+	jsStr := toolCallRE.FindString(text)
+	if jsStr == "" {
+		return ""
+	}
+	jsStr = strings.TrimSpace(jsStr)
+	jsStr = strings.TrimPrefix(jsStr, "__tool_call__")
+	jsStr = strings.TrimSuffix(jsStr, "__tool_call__")
+	jsStr = strings.TrimSpace(jsStr)
+
+	start := strings.Index(jsStr, "{")
+	end := strings.LastIndex(jsStr, "}")
+	if start == -1 || end == -1 || end <= start {
+		return ""
+	}
+	jsStr = jsStr[start : end+1]
+
+	var fc models.FuncCall
+	if err := json.Unmarshal([]byte(jsStr), &fc); err != nil {
+		return ""
+	}
+	return fc.Name
+}