feat: 98 generative AI agent for classification

DEMO/genAI_prototype.ipynb

@@ -0,0 +1,262 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# GEN AI Workflow Notebook ✨✨"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "ClassifAI allows users to create vector databases from text datasets, search those vector datasets in an ad-hoc retrieval manner, and deploy this pipeline as a restAPI service using FastAPI to perform AI assisted classification tasks.\n",
+    "\n",
+    "\n",
+    "An recent emergent AI field is Retrieval Augmented Generation (RAG) where text generation models that traditionally respond to a user 'prompt', first retrieve relevant infomration from a vector database via adhoc retrieval processes, and use those serach results as context to generate an answer for the original user prompt.\n",
+    "\n",
+    "#### This notebook shows how we use RAG agents to perform classification on our VectorStore semantic search results!\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## ClassifAIs existing retrieval setup\n",
+    "\n",
+    "![Server_Image](files/servers.png)\n",
+    "\n",
+    "#### The other modules of ClassifAI provide 3 core classes that work together to:\n",
+    "\n",
+    "1. Vectorisers, to create embeddings,\n",
+    "2. Vectorstores, to create datgabsees of vectors and the ability to searc/query those database\n",
+    "3. RestAPI, to deploy a rest api service on a server to allow connections that search the created vector databases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!uv pip install \"classifai[gcp]\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "We need to first set up a traditional ClassifAI pipeline that can provide search/classification results for our generative model to use as context...\n",
+    "\n",
+    "\n",
+    "All of the following code is part of our standard ClassifAI setupm, and is a short demo of how you can create a semantic search classification system. <b>Check out our general_workflow_demo.ipynb notebook for a walkthrough of the content of these cells!</b>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!gcloud auth application-default login"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from classifai.indexers import VectorStore\n",
+    "from classifai.vectorisers import HuggingFaceVectoriser\n",
+    "\n",
+    "# Our embedding model is pulled down from HuggingFace, or used straight away if previously downloaded\n",
+    "# This also works with many different huggingface models!\n",
+    "vectoriser = HuggingFaceVectoriser(model_name=\"sentence-transformers/all-MiniLM-L6-v2\")\n",
+    "\n",
+    "\n",
+    "my_vector_store = VectorStore(\n",
+    "    file_name=\"./data/fake_soc_dataset.csv\",  # demo csv file from the classifai package repo! (try some of our other DEMO/data datasets)\n",
+    "    data_type=\"csv\",\n",
+    "    vectoriser=vectoriser,\n",
+    "    agent=None,\n",
+    "    overwrite=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## We've set up our semantic search classification system\n",
+    "\n",
+    "The following cell runs the 'search' method of the vectorstore, which will take a query for our created vectorstore and return the top k (n_results) most similar samples stored in the vectorstore. This is an example of our exisiting retrieval capabilities with the ClassifAI package. \n",
+    "\n",
+    "This retrieved set can then be used to make a final classification decision, by some method such as:\n",
+    "- automatically choosing the top ranked item, \n",
+    "- a human in the loop physically looking at the retrieved candidates making the decision,\n",
+    "- by using a generative AI agent to assess the candidate lists and making a final decison... more on that later"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from classifai.indexers.dataclasses import VectorStoreSearchInput\n",
+    "\n",
+    "input_object = VectorStoreSearchInput(\n",
+    "    {\"id\": [1, 2], \"query\": [\"dairy famer that sells milk\", \"happy construction worker\"]}\n",
+    ")\n",
+    "\n",
+    "bb = my_vector_store.search(input_object, n_results=5)\n",
+    "bb"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Creating an AI Agent\n",
+    "\n",
+    "With our semantic search classification pipeline set up, we can send the top K results (seen above) to a Genereative AI model and ask that LLM to make a final decision on which result is the correct result for the user's original input query. Passing semantic search results to a generatieve model for some task is often referred to as 'Retrieval Augmented Generation'.\n",
+    "\n",
+    "\n",
+    "![Rag_Image](files/agent.png) \n",
+    "\n",
+    "\n",
+    "Our `GcpAgent` class agent has a transform() method which accepts a `VectorStoreSearchOutput` object. Passing this `VectorStoreSearchOutput` object to the transform() method will return another `VectorStoreSearchOutput` object, which will modify the results in some way. The classificaion GcpAgent, reduces the semantic search results down to a single result row for each query in the VectorStoreSearchOutput results.\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To instantiate the GcpAgent, the constructor takes:\n",
+    "\n",
+    "- project_id - of the google project associated with the Google Gemini embedding models\n",
+    "- location - corresponding to the Gcloud Project\n",
+    "- model_name - the specific LLM model to use, that is avaiable on Gcloud\n",
+    "- task_type - indicates the kind of work you want the LLM to do on the ClassifAI results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from classifai.agents import GcpAgent\n",
+    "\n",
+    "my_agent = GcpAgent(\n",
+    "    project_id=\"xxxxxx\", location=\"europe-west2\", model_name=\"gemini-2.5-flash\", task_type=\"classification\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can pass some VectorStore search results directly to the agents transform() method"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "semantic_search_results = my_vector_store.search(input_object, n_results=5)\n",
+    "\n",
+    "print(type(semantic_search_results))\n",
+    "semantic_search_results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### To use the If the classification agent can make a decision, it will return a single row for the corresponding query_id. Otherwise it will return the original results\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_agent.transform(semantic_search_results)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### To use the agent in our ClassifAI pipeline, we can attach it to the running VectorStore"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_vector_store.agent = my_agent"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### We can then run vectorstore.search and it will automatically call the agent to process the results!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_vector_store.search(input_object, n_results=5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Thats it!\n",
+    "\n",
+    "You can see in the above cell that my_agent.transform() takes and returns a VectorStoreSearchOutput object and therefore integrates with the VectorStore search method.\n",
+    "\n",
+    "Check out the general_worfklow notebook to see how VectorStores can be deployed, and you'll find that the Agent can be deployed as an integrated part of the VectorStore."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/classifai/agents/__init__.py

@@ -0,0 +1,7 @@
+from .base import AgentBase
+from .gcp import GcpAgent
+
+__all__ = [
+    "AgentBase",
+    "GcpAgent",
+]

src/classifai/agents/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from classifai.indexers.dataclasses import VectorStoreSearchOutput
+
+##
+# The following is the abstract base class for all RAG generative models.
+##
+
+
+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