Routing Queries to Data Sources

1
import json
2
import os
3
import cohere
4
5
from tool_def import (
6
    search_developer_docs,
7
    search_developer_docs_tool,
8
    search_internet,
9
    search_internet_tool,
10
    search_code_examples,
11
    search_code_examples_tool,
12
)
13
14
co = cohere.ClientV2("COHERE_API_KEY") # Get your free API key: https://dashboard.cohere.com/api-keys

1
from langchain_community.tools.tavily_search import TavilySearchResults
2
3
def search_developer_docs(query: str) -> dict:
4
    
5
    developer_docs = [{"text" : "## The Rerank endpoint\nThis endpoint takes in a query and a list of texts and produces an ordered array with each text assigned a relevance score."},
6
    
7
    {"text" : "## The Embed endpoint\nThis endpoint returns text embeddings. An embedding is a list of floating point numbers that captures semantic information about the text that it represents.."},
8
9
    {"text" : "## Embed endpoint multilingual support\nIn addition to embed-english-v3.0 we offer a best-in-class multilingual model embed-multilingual-v3.0 with support for over 100 languages."},
10
    
11
    {"text" : "## The Chat endpoint\nThis endpoint facilitates a conversational interface, allowing users to send messages to the model and receive text responses."},
12
    
13
    {"text" : "## Retrieval Augmented Generation (RAG)\nRAG is a method for generating text using additional information fetched from an external data source, which can greatly increase the accuracy of the response."},
14
    
15
    {"text" : "## The temperature parameter\nTemperature is a number used to tune the degree of randomness of a generated text."},
16
    ]
17
18
    return developer_docs
19
20
def search_internet(query: str) -> dict:
21
    tool = TavilySearchResults(
22
        max_results=5,
23
        search_depth="advanced",
24
        include_answer=True,
25
        include_raw_content=True
26
    )
27
    documents = tool.invoke({"query": query})
28
    
29
    return documents
30
31
def search_code_examples(query: str) -> dict:
32
    
33
    code_examples = [
34
        {"content": "Calendar Agent with Native Multi Step Tool"},
35
        {"content": "Wikipedia Semantic Search with Cohere Embedding Archives"},
36
        {"content": "RAG With Chat Embed and Rerank via Pinecone"},
37
        {"content": "Build Chatbots That Know Your Business with MongoDB and Cohere"},
38
        {"content": "Advanced Document Parsing For Enterprises"}
39
    ]
40
    
41
    return code_examples
42
43
functions_map = {
44
    "search_developer_docs": search_developer_docs,
45
    "search_internet": search_internet,
46
    "search_code_examples": search_code_examples,
47
}

1
search_developer_docs_tool = {
2
        "type": "function",
3
        "function": {
4
            "name": "search_developer_docs",
5
            "description": "Searches the Cohere developer documentation. Use this tool for queries related to the Cohere API, SDKs, or other developer resources.",
6
            "parameters": {
7
                "type": "object",
8
                "properties": {
9
                    "query": {
10
                        "type": "string",
11
                        "description": "The search query."
12
                    }
13
                },
14
                "required": ["query"]
15
            }
16
        }
17
}
18
    
19
search_internet_tool = {
20
        "type": "function",
21
        "function": {
22
            "name": "search_internet",
23
            "description": "Searches the internet. Use this tool for general queries that would not be found in the developer documentation.",
24
            "parameters": {
25
                "type": "object",
26
                "properties": {
27
                    "query": {
28
                        "type": "string",
29
                        "description": "The search query."
30
                    }
31
                },
32
                "required": ["query"]
33
            }
34
        }
35
}
36
37
search_code_examples_tool = {
38
        "type": "function",
39
        "function": {
40
            "name": "search_code_examples",
41
            "description": "Searches code examples and tutorials on using Cohere.",
42
            "parameters": {
43
                "type": "object",
44
                "properties": {
45
                    "query": {
46
                        "type": "string",
47
                        "description": "The search query."
48
                    }
49
                },
50
                "required": ["query"]
51
            }
52
        }
53
}

1
tools = [
2
    search_developer_docs_tool,
3
    search_internet_tool,
4
    search_code_examples_tool
5
]

1
system_message="""## Task and Context
2
You are an assistant who helps developers use Cohere. You are equipped with a number of tools that can provide different types of information. If you can't find the information you need from one tool, you should try other tools if there is a possibility that they could provide the information you need."""

1
model = "command-r-plus-08-2024"
2
3
def run_agent(query, messages=None):
4
    if messages is None:
5
        messages = []
6
        
7
    if "system" not in {m.get("role") for m in messages}:
8
        messages.append({"role": "system", "content": system_message})
9
    
10
    # Step 1: get user message
11
    print(f"QUESTION:\n{query}")
12
    print("="*50)
13
    
14
    messages.append({"role": "user", "content": query})
15
16
    # Step 2: Generate tool calls (if any)
17
    response = co.chat(
18
        model=model,
19
        messages=messages,
20
        tools=tools,
21
        temperature=0.1
22
    )
23
24
    while response.message.tool_calls:
25
        
26
        print("TOOL PLAN:")
27
        print(response.message.tool_plan,"\n")
28
        print("TOOL CALLS:")
29
        for tc in response.message.tool_calls:
30
            print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
31
        print("="*50)
32
33
        messages.append({"role": "assistant", "tool_calls": response.message.tool_calls, "tool_plan": response.message.tool_plan})        
34
        
35
        # Step 3: Get tool results
36
        for tc in response.message.tool_calls:
37
            tool_result = functions_map[tc.function.name](**json.loads(tc.function.arguments))
38
            tool_content = []
39
            for data in tool_result:
40
                tool_content.append({"type": "document", "document": {"data": json.dumps(data)}})
41
                # Optional: add an "id" field in the "document" object, otherwise IDs are auto-generated
42
            messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_content})
43
        
44
        # Step 4: Generate response and citations 
45
        response = co.chat(
46
            model=model,
47
            messages=messages,
48
            tools=tools,
49
            temperature=0.1
50
        )
51
    
52
    messages.append({"role": "assistant", "content": response.message.content[0].text})
53
        
54
    # Print final response
55
    print("RESPONSE:")
56
    print(response.message.content[0].text)
57
    print("="*50)
58
    
59
    # Print citations (if any)
60
    verbose_source = False # Change to True to display the contents of a source
61
    if response.message.citations:
62
        print("CITATIONS:\n")
63
        for citation in response.message.citations:
64
            print(f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' ")
65
            print("Sources:")
66
            for idx, source in enumerate(citation.sources):
67
                print(f"{idx+1}. {source.id}")
68
                if verbose_source:
69
                    print(f"{source.tool_output}")
70
            print("\n") 
71
                    
72
    return messages

1
messages = run_agent("How many languages does Embed support?")

1
messages = run_agent("Who are the authors of the sentence BERT paper?")

1
messages = run_agent("Any tutorials that are relevant for enterprises?")