How do i create tools that need credentials?

[Nathan-Intergral]

Hi all,

I was wondering if someone could help me out that something that has been a sore point for a while in my application.

I started out making my autogen application following the docs for adding tools the standard way, using the register_for_llm etc etc. However quickly ran in to a problem when i needed to provide auth credentials for an account in the tool for an api call. I ended up trying to solve this with a singleton to store these details but realised quickly that wouldn't work as am working with a service that accepts requests from multiple tenants at any given time.

I ended up switching to using functions and using Langchain Tool classes so that i could initialise the class with any parameters i wanted, auth token etc. Which is working for me now.

However, i have recently learned that functions are a deprecated feature of OpenAI and theres several downsides to using functions over tools. With that in mind, id like to switch back to tools but I have no idea how to get around the issue i described above.

Does anyone know how i can provide parameters like auth headers, contexts etc to tools in a multi tenant environment?

Thanks in advance.

[ekzhu]

I ended up switching to using functions and using Langchain Tool classes so that i could initialise the class with any parameters i wanted, auth token etc. Which is working for me now.

Could you provide an example of how this solves your problem? If you want to have different credentials for different instance of invocation, you can use a simple function wrapper.

[Nathan-Intergral]

Yeah sure thing. Here's one of my tools currently, I pass it the user credentials for the authed user that sent the request and the context object which stores conversation context management stuff i need to check the state of things throughout the duration of the conversation.

class ValidateQueryToolInput(BaseModel):
    query: Annotated[str, "the query you want to validate."]


class ValidateQueryTool(BaseTool):
    name = "ValidateQueryTool"
    description = "Use this to validate a generated query is valid promQL syntax."
    user: auth_schemas.User = None
    context: ContextManager = None
    args_schema: Type[BaseModel] = ValidateQueryToolInput

    def __init__(self, user: auth_schemas.User, context: ContextManager):
        super().__init__()
        self.user = user
        self.context = context

    def _run(self, input: str):
        return "Does not support sync"

    async def _arun(self, query: str):
        url = settings.prometheus_url + "/api/v1/query"
        timeout = httpx.Timeout(10.0, read=15.0)
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.get(
                url,
                headers={"X-Scope-OrgID": self.user.account_id},
                params={"query": query},
            )

            if response.status_code == 400:
                try:
                    data = response.json()
                    error = data.get("error")
                    return str(error) if error else "invalid promQL"
                except Exception as e:
                    logger.warning(
                        f"Unable to get prometheus response for query {query}."
                        f" Caught exception: {e}"
                    )
                    return (
                        f"Unable to get prometheus response for query {query}."
                        f" Caught exception: {e}"
                    )
            elif response.status_code != 200:
                return (
                    f"Got status code {response.status_code} when attempting to"
                    f" validate query: {query}"
                )

        try:
            self.context.update_query_params(query_string=query, validated=True)
        except QueryDoesNotExistError:
            return (
                "promQL is valid but the query has not been created, please create the"
                " query then revalidate it"
            )

        return f"promQL is valid for query {query}"

Thanks for taking a look

[ekzhu]

Not sure if this helps:

# Define your tool
def validate_query(
    user: auth_schemas.User,
    context: ContextManager,
    query: str,
) -> str:
    # your implementation
    pass

# in your application code
user = ...
context = ...

def validate_query_tool(query):
    return validate_query(user, context, query)

functools.update_wrapper(wrapper, validate_query)

# Now your tool is ready to use.

This is just a quick sketch, you can take a look at functools.update_wrapper for more details -- you might want to tweak a bit.

[Nathan-Intergral]

ah interesting.. Okay i get the idea, ill take a look. Thank you!

[Abhishektght]

how can we use this with autogen assistant agent after registering it?