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