feat: 98 generative AI agent for classification

[frayle-ons]

✨ Summary

These changes introduce the framework for using Generative AI Large Language models to perform classifications on VectorStoreSearchResult objects, and integrates with the existing VectorStore search pipeline.

It is adding a new module called agents, which function similarly to the Vectorisers class/module. Currently implemented is the GcpVectoriser. This class has a transform() method which intakes a VectorStoreSearchOutput object and transforms it in some way using generative AI but still outputs a valid VectorStoreSearchOutput object.

For example, we can instantiate an agent that classifies the top K semantic search candidates with the following Python code:

from classifai.agents import GcpAgent

my_agent = GcpAgent(
    project_id="xxxxxx", 
    location="europe-west2", 
    model_name="gemini-2.5-flash", 
    task_type="classification"
)

The instantiated agent has a transform method that accepts a VectorStoreSearchOutput object:

my_agent.transform(semantic_search_results)

This standalone agent can then be injected into the VectorStore, so that it automatically runs when the VectorStore.search()method is called, by setting the 'agent' attribute. Either on instantiation:

my_vector_store = VectorStore(
    file_name="./data/fake_soc_dataset.csv",  # demo csv file from the classifai package repo! (try some of our other DEMO/data datasets)
    data_type="csv",
    vectoriser=vectoriser,
    agent=my_agent
    overwrite=True,
)

or by setting the attribute to a VectorStore already in runtime: my_vector_store.agent = my_agent

Inner workings and framework design

The GcpVectoriser inherits from a base class that specifies the Agents must take in and output a VectorStoreSearchOutput object.

class Agentase(ABC):
    """Abstract base class for all Generative and RAG models."""

    @abstractmethod
    def transform(
        self,
        results: VectorStoreSearchOutput,
    ) -> VectorStoreSearchOutput:
        """Passes VectorStoreSearchOutput object, which the Agent manipulates in some way and returns."""
        pass

This introduces a core concepts of how we use Generative AI in the ClassifAI workflow - i.e. it should only manipulate existing results objects.

There are 4 key components to the GcpAgent model that has been created:

1. The client API that connects to the Google cloud service, where we can actually send text to a generative model and get a response,
2. A function that formats Vector Store search results into a more natural text format for the generative model to read.
3. A system prompt that describes to the generative model what to do with the vector store results
4. A post-processing function that validates the generative model's output and performs any corresponding operations on the VectorStoreSearchOutput object.

using this 4 step approach makes it possible to create different ways of using the Agent, by changing the system prompt instruction and the corresponding post-processing function - all while grounding the behaviour to make sure we're only manipulating valid VectorStoreSearchOutput objects.

Other agent models such as a HuggingfaceAgent (still to be developed), could follow the same paradigm, and it would give guidance to package users who may wish to implement their own custom generative model by inheriting from the base class.

Classification example

Both the System Prompt and the post-processing function can vary in behaviour based on which task type the agent model is instantiated with. Currently only 'classification' task type is available.

With the classification task type the system prompt is set as (paraphrasing for briefness):

CLASSIFICATION SYSTEM PROMPT = """
You are a classification model that will be provided a user input query and 5 candidate semantic 
search results, with corresponding IDs, from 0 to 4. You must choose the ID of the candidate 
result you think best matches the user's input text. Output your answer in the following format, 
and output -1 if the result is unclear:

{
classification : <your chosen ID>
}

"""

(the actual system prompt is more detailed, contains an example of the data structure, and guidelines - see this in the code files)

The system prompt and formatted search results are combined and passed to the generative model.

The classification post-processing function then assesses if the generative AI output is of the correct format using Pydantic. If it is, it takes the chosen ID and reduces the original VectorStoreSearchOutput object down to the chosen row using the ID generated by the generative AI. If the correct format response was not generated, the post-processing function simply returns the original VectorStoreSearchOutput with no changes.

Considerations:

• Hardcoded the number of semantic search classification results the generative model uses to 5 - this could be changed to a settable parameter, but hardcoded for now due to larger value of N increasing the cost of using models like this
• When a batch of search results are passed to the model, it will pass each query result to the generative API sequentially. The batch api on Cloud is meant to be used for longer timescales and is not appropriate for batch inference in this way. We could introduce asynchronous calls to the cloud api instead, that way making multiple API calls at once but not added at this time.
• The system prompt instructs the model to output a value of -1 when it is unsure of the correct classification. This is somewhat ambiguous as there is a different between uncodable and 'cannot code for other reasons'. introducing different reasons that a decision can't be made would complicate the system prompt.

Agents as hooks instead of as a specific attribute of the VectorStore

One specific consideration, since all the agent is doing is transforming the VectorStoreSearchOutput object into another VectorStoreSearchOutput: it could be treated simply as a post-processing hook on the VectorStore search method. This would be in line with the hooks update we did recently and would simplify the integration of the agent with the VectorStore, although the integration currently is not complex. We could just provide the agents as a kind of fancy pre-build hook ready for use by users.

📜 Changes Introduced

• Introduces new module and base class for 'agents'
• Introduces the GCPAgent class for classification of VectorStoreSearchOutput objects
• Currated system prompt and post-processing function to handle generative LLM input and output
• Integration of Agent with VectorStore so that the agent operates automatically on vector store search results
• New notebook that showcases setting up the VectorStore and Agent to do classification

✅ Checklist

Please confirm you've completed these checks before requesting a review.

• Code passes linting with Ruff
• Security checks pass using Bandit
• API and Unit tests are written and pass using pytest
• Terraform files (if applicable) follow best practices and have been validated (terraform fmt & terraform validate)
• DocStrings follow Google-style and are added as per Pylint recommendations
• Documentation has been updated if needed

🔍 How to Test

Installing this branch as normal, and then running through the new notebook would be a good way to explore and test the features of the changes. One note, is to make sure to use a google cloud project with the generative language api activated. Otherwise the requests to that api won't be successful

[lukeroantreeONS]

I've added some comments for the current implementation, but I also think it's worth doing a PoC (maybe in another branch) of a post-processing 'hook' type implementation so we can compare the pros/cons of the two approaches.

[lukeroantreeONS]

Spelling

[lukeroantreeONS]

If we're going with a solid XML structure requirement, should we enforce (some) input sanitisation to ensure user-input doesn't break the XML? (e.g. escape '<')

[lukeroantreeONS]

Similar concerns with other characters if we opt for a JSON-like approach instead etc.

[lukeroantreeONS]

My preference would be to use a dedicated tool for forming the XML, to avoid surprises / edge cases / etc. (e.g. https://pypi.org/project/dicttoxml/)

[lukeroantreeONS]

Do we want to lock this restriction in rather than let the user specify a max?

[lukeroantreeONS]

Do we want to enforce this range for all use-cases? I think making it user-configurable would be good. Happy to be convinced otherwise though

[lukeroantreeONS]

Preference would be to have a class-level list of supported options, which is used to format this error message - centralises where these are defined, helping avoid missing updating the error etc. if we add options down the line.

[lukeroantreeONS]

It would be good to think early about making this asynchronous - each generative request will be slow, but with little to do on this side while waiting

[matweldon]

Suggested change

	class Agentase(ABC):
	class AgentBase(ABC):

[matweldon]

I like this as a demonstration of what's possible with the system. In general though, it feels too specific and 'brittle' to become a fixed part of the package, if you know what I mean? For example, the prompt hard codes the number of results retrieved. Would a more general system have the user constructing a templated prompt, rather than us constructing a templated prompt for the user? What other scaffolding could we give users in constructing RAG post-processors without it becoming too brittle?

I'm more convinced now that this would be better implemented as a post-processing hook, because then we could present it as more of an example and less of a fixed feature?