examples/basic/fn-call-local-simple.py - main - githubmirror/langroid

examples/basic/fn-call-local-simple.py
156 строк5 KB
Prasad Chalasani
fix: relevance_extractor_agent improve tool prompt
30 сен 2025, 04:50
30 сен 2025, 04:50500fa88
"""
Function-calling example using a local/remote open LLM.

"Function-calling" refers to the ability of the LLM to generate
a structured response, typically a JSON object, instead of a plain text response,
which is then interpreted by your code to perform some action.
This is also referred to in various scenarios as "Tools", "Actions" or "Plugins".
See more here: https://langroid.github.io/langroid/quick-start/chat-agent-tool/

Run like this (to run with llama-3.1-8b-instant via groq):

python3 examples/basic/fn-call-local-simple.py -m groq/llama-3.1-8b-instant

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


"""

import os
from typing import List

import fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_document import ChatDocument
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import FinalResultTool
from pydantic import BaseModel, Field
from langroid.utils.configuration import settings

# for best results:
DEFAULT_LLM = lm.OpenAIChatModel.GPT4o

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# (1) Define the desired structure via Pydantic.
# Here we define a nested structure for City information.
# The "Field" annotations are optional, and are included in the system message
# if provided, and help with generation accuracy.


class CityData(BaseModel):
    population: int = Field(..., description="population of city")
    country: str = Field(..., description="country of city")


class City(BaseModel):
    name: str = Field(..., description="name of city")
    details: CityData = Field(..., description="details of city")


# (2) Define the Tool class for the LLM to use, to produce the above structure.
class CityTool(lr.agent.ToolMessage):
    """Present information about a city"""

    request: str = "city_tool"
    purpose: str = """
    To present <city_info> AFTER user gives a city name,
    with all fields of the appropriate type filled out;
    """
    city_info: City = Field(..., description="information about a city")

    def handle(self) -> FinalResultTool:
        """Handle LLM's structured output if it matches City structure"""
        print("SUCCESS! Got Valid City Info")
        return FinalResultTool(answer=self.city_info)

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        # Used to provide few-shot examples in the system prompt
        return [
            cls(
                city_info=City(
                    name="San Francisco",
                    details=CityData(
                        population=800_000,
                        country="USA",
                    ),
                )
            )
        ]


def app(
    m: str = DEFAULT_LLM,  # model
    d: bool = False,  # pass -d to enable debug mode (see prompts etc)
    nc: bool = False,  # pass -nc to disable cache-retrieval (i.e. get fresh answers)
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=32000,  # set this based on model
        max_output_tokens=1000,
        temperature=0.2,
        stream=True,
        timeout=45,
    )

    # Recommended: First test if basic chat works with this llm setup as below:
    # Once this works, then you can try the rest of the example.
    #
    # agent = lr.ChatAgent(
    #     lr.ChatAgentConfig(
    #         llm=llm_cfg,
    #     )
    # )
    #
    # agent.llm_response("What is 3 + 4?")
    #
    # task = lr.Task(agent)
    # verify you can interact with this in a chat loop on cmd line:
    # task.run("Concisely answer some questions")

    # Define a ChatAgentConfig and ChatAgent

    config = lr.ChatAgentConfig(
        llm=llm_cfg,
        handle_llm_no_tool=f"""
            You FORGOT to use the TOOL/Function `{CityTool.name()}` 
            to present city info!
            """,
        system_message=f"""
        You will receive a city name, 
        and you must use the TOOL/FUNCTION `{CityTool.name()}` to generate/present
        information about the city. In other words, your response must 
        be a JSON string starting with `{{"request": "{CityTool.name()}", ...}}`
        """,
    )

    agent = lr.ChatAgent(config)

    # (4) Enable the Tool for this agent --> this auto-inserts JSON instructions
    # and few-shot examples (specified in the tool defn above) into the system message
    agent.enable_message(CityTool)

    # (5) Create task specialized to return City object
    task: City | None = lr.Task(agent, interactive=False)[City]

    while True:
        city = Prompt.ask("Enter a city name")
        if city in ["q", "x"]:
            break
        result: City | None = task.run(city)
        if result:
            print(f"City Info: {result}")
        else:
            print("No valid city info found.")


if __name__ == "__main__":
    fire.Fire(app)