diff options
Diffstat (limited to 'tools.go')
| -rw-r--r-- | tools.go | 77 |
1 files changed, 75 insertions, 2 deletions
@@ -16,6 +16,7 @@ import ( "sync" "time" + "gf-lt/rag" "github.com/GrailFinder/searchagent/searcher" ) @@ -58,9 +59,9 @@ Your current tools: "when_to_use": "when asked to search the web for information; returns clean summary without html,css and other web elements; limit is optional (default 3)" }, { -"name":"websearch_raw", +"name":"rag_search", "args": ["query", "limit"], -"when_to_use": "when asked to search the web for information; returns raw data as is without processing; limit is optional (default 3)" +"when_to_use": "when asked to search the local document database for information; performs query refinement, semantic search, reranking, and synthesis; returns clean summary with sources; limit is optional (default 3)" }, { "name":"read_url", @@ -146,6 +147,7 @@ under the topic: Adam's number is stored: After that you are free to respond to the user. ` webSearchSysPrompt = `Summarize the web search results, extracting key information and presenting a concise answer. Provide sources and URLs where relevant.` + ragSearchSysPrompt = `Synthesize the document search results, extracting key information and presenting a concise answer. Provide sources and document IDs where relevant.` readURLSysPrompt = `Extract and summarize the content from the webpage. Provide key information, main points, and any relevant details.` summarySysPrompt = `Please provide a concise summary of the following conversation. Focus on key points, decisions, and actions. Provide only the summary, no additional commentary.` basicCard = &models.CharCard{ @@ -170,6 +172,10 @@ func init() { panic("failed to init seachagent; error: " + err.Error()) } WebSearcher = sa + + if err := rag.Init(cfg, logger, store); err != nil { + logger.Warn("failed to init rag; rag_search tool will not be available", "error", err) + } } // getWebAgentClient returns a singleton AgentClient for web agents. @@ -196,6 +202,8 @@ func getWebAgentClient() *agent.AgentClient { func registerWebAgents() { webAgentsOnce.Do(func() { client := getWebAgentClient() + // Register rag_search agent + agent.Register("rag_search", agent.NewWebAgentB(client, ragSearchSysPrompt)) // Register websearch agent agent.Register("websearch", agent.NewWebAgentB(client, webSearchSysPrompt)) // Register read_url agent @@ -239,6 +247,48 @@ func websearch(args map[string]string) []byte { return data } +// rag search (searches local document database) +func ragsearch(args map[string]string) []byte { + query, ok := args["query"] + if !ok || query == "" { + msg := "query not provided to rag_search tool" + logger.Error(msg) + return []byte(msg) + } + limitS, ok := args["limit"] + if !ok || limitS == "" { + limitS = "3" + } + limit, err := strconv.Atoi(limitS) + if err != nil || limit == 0 { + logger.Warn("ragsearch limit; passed bad value; setting to default (3)", + "limit_arg", limitS, "error", err) + limit = 3 + } + + ragInstance := rag.GetInstance() + if ragInstance == nil { + msg := "rag not initialized; rag_search tool is not available" + logger.Error(msg) + return []byte(msg) + } + + results, err := ragInstance.Search(query, limit) + if err != nil { + msg := "rag search failed; error: " + err.Error() + logger.Error(msg) + return []byte(msg) + } + + data, err := json.Marshal(results) + if err != nil { + msg := "failed to marshal rag search result; error: " + err.Error() + logger.Error(msg) + return []byte(msg) + } + return data +} + // web search raw (returns raw data without processing) func websearchRaw(args map[string]string) []byte { // make http request return bytes @@ -997,6 +1047,7 @@ var fnMap = map[string]fnSig{ "recall": recall, "recall_topics": recallTopics, "memorise": memorise, + "rag_search": ragsearch, "websearch": websearch, "websearch_raw": websearchRaw, "read_url": readURL, @@ -1033,6 +1084,28 @@ func callToolWithAgent(name string, args map[string]string) []byte { // openai style def var baseTools = []models.Tool{ + // rag_search + models.Tool{ + Type: "function", + Function: models.ToolFunc{ + Name: "rag_search", + Description: "Search local document database given query, limit of sources (default 3). Performs query refinement, semantic search, reranking, and synthesis.", + Parameters: models.ToolFuncParams{ + Type: "object", + Required: []string{"query", "limit"}, + Properties: map[string]models.ToolArgProps{ + "query": models.ToolArgProps{ + Type: "string", + Description: "search query", + }, + "limit": models.ToolArgProps{ + Type: "string", + Description: "limit of the document results", + }, + }, + }, + }, + }, // websearch models.Tool{ Type: "function", |