Skip to content

LLM - Retrieval

Install packages

1
2
3
4
!uv pip install -q \
    litellm==1.78.5 \
    python-dotenv==1.1.1 \
    pydantic==2.12.3

Import packages

1
2
3
4
5
6
7
import json

import litellm
from dotenv import load_dotenv
from pydantic import BaseModel, Field

load_dotenv()

True

Define knowledge base

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
def search_kb(question: str):
    return str(
        {
            "records": [
                {
                    "id": 1,
                    "question": "What is the return policy?",
                    "answer": "Items can be returned within 30 days of purchase with original receipt. Refunds will be processed to the original payment method within 5-7 business days.",
                },
                {
                    "id": 2,
                    "question": "Do you ship internationally?",
                    "answer": "Yes, we ship to over 50 countries worldwide. International shipping typically takes 7-14 business days and costs vary by destination. Please note that customs fees may apply.",
                },
                {
                    "id": 3,
                    "question": "What payment methods do you accept?",
                    "answer": "We accept Visa, Mastercard, American Express, PayPal, and Apple Pay. All payments are processed securely through our encrypted payment system.",
                },
            ]
        }
    )

Define tools specification

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
tools = [
    {
        "type": "function",
        "name": "search_kb",
        "description": "Get the answer to the user's question from the knowledge base.",
        "parameters": {
            "type": "object",
            "properties": {
                "question": {"type": "string"},
            },
            "required": ["question"],
            "additionalProperties": False,
        },
        "strict": True,
    }
]

Define messages

  • System message, defined once
  • User messages
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
messages = [
    {
        "role": "system",
        "content": "You are a helpful assistant that answers questions from the knowledge base about our e-commerce store.",
    },
    {
        "role": "user",
        "content": "What is the return policy?",
    },
]

Get chat completion

1
2
3
4
5
6
7
completion = litellm.completion(
    model="gemini/gemini-2.0-flash",
    messages=messages,
    tools=tools,
)

print(completion.choices[0].message.tool_calls)

[ChatCompletionMessageToolCall(index=0, function=Function(arguments='{"question": "What is the return policy?"}', name='search_kb'), id='call_2ed16bb9ed424be8aeb1eea38408', type='function')]

Call the function and append its result to messages list

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
def call_function(name, args):
    if name == "search_kb":
        return search_kb(**args)


for tool_call in completion.choices[0].message.tool_calls:
    name = tool_call.function.name
    args = json.loads(tool_call.function.arguments)
    messages.append(completion.choices[0].message)

    result = call_function(name, args)
    messages.append(
        {
            "role": "tool",
            "tool_call_id": tool_call.id,
            "content": json.dumps(result),
        }
    )

messages

[{'role': 'system',
  'content': 'You are a helpful assistant that answers questions from the knowledge base about our e-commerce store.'},
 {'role': 'user', 'content': 'What is the return policy?'},
 Message(content=None, role='assistant', tool_calls=[ChatCompletionMessageToolCall(index=0, function=Function(arguments='{"question": "What is the return policy?"}', name='search_kb'), id='call_2ed16bb9ed424be8aeb1eea38408', type='function')], function_call=None, images=[], thinking_blocks=[], provider_specific_fields=None),
 {'role': 'tool',
  'tool_call_id': 'call_2ed16bb9ed424be8aeb1eea38408',
  'content': '"{\'records\': [{\'id\': 1, \'question\': \'What is the return policy?\', \'answer\': \'Items can be returned within 30 days of purchase with original receipt. Refunds will be processed to the original payment method within 5-7 business days.\'}, {\'id\': 2, \'question\': \'Do you ship internationally?\', \'answer\': \'Yes, we ship to over 50 countries worldwide. International shipping typically takes 7-14 business days and costs vary by destination. Please note that customs fees may apply.\'}, {\'id\': 3, \'question\': \'What payment methods do you accept?\', \'answer\': \'We accept Visa, Mastercard, American Express, PayPal, and Apple Pay. All payments are processed securely through our encrypted payment system.\'}]}"'}]

Define output structure

1
2
3
4
5
6
7
class KBResponse(BaseModel):
    answer: str = Field(
        description="The answer to the user's question from the knowledge base."
    )
    source: str = Field(
        description="The record id of the knowledge base entry that was used to answer the question."
    )

Supply results and call the model again

1
2
3
4
5
6
7
8
completion_2 = litellm.completion(
    model="gemini/gemini-2.0-flash",
    messages=messages,
    tools=tools,
    response_format=KBResponse,
)

print(completion_2.choices[0].message.content)

{

  "answer": "Items can be returned within 30 days of purchase with original receipt. Refunds will be processed to the original payment method within 5-7 business days.",

  "source": "1"

}