Generative Agentsを作成するAPIサーバと管理画面をローカルで作って遊んでみた

はじめに

こんにちは、エンジニアのクロ(@kro96_xr)です。

先月の頭にGenerative Agentsに関する論文が話題になりましたね。

arxiv.org

そして、このGenerative AgentsについてはLangChainで実装例が提供されています。

python.langchain.com

この機能をWebAPI化すれば管理画面でエージェントを作成、他のクライアントから呼び出すこともできるのでは?と考えて実装してみました。
本記事では主にWebAPIの実装部分について触れていきます。

実際には、データや記憶の永続化がうまくできていないのでサーバ再起動で消えてしまう状態ですが、ローカルでちょっと遊ぶ分には十分だと思います。

構成

構成(技術スタック)は以下のようになります。フロントエンドは専門ではないので雰囲気で実装しています。

  • フロントエンド

    • フレームワーク:React
    • デザイン:MUI
  • バックエンド

    • 実行環境:Anaconda on Docker
    • フレームワーク:FastAPI

今回はDocker上でAnacondaを動かしていますが、これはFAISSのインストール上必要だったためです。

github.com

実装後イメージ

LangChainのページのサンプルと同様にエージェントにインタビューしてみるとこのようになります。

プロンプトを調整してやると恋愛ゲームのようにすることもできます。

qiita.com

ChatGPTにキャラクター設定を作ってもらい…

エージェントを生成するといい感じに思い出をサマライズしてくれます。

チャットもそれっぽくなります。断られましたが…

実装

それでは具体的な実装をみていきます。

ディレクトリ構成や実際のコードはGithubのリポジトリをご覧ください。

github.com

まずはコンテナの構築からです。

compose.yaml

version: "3.9"

services:
  agent:
    build: ./ai-agent-server
    ports:
      - "5001:5001"
    volumes:
      - ./ai-agent-server/app:/app
    tty: true
    environment:
      OPENAI_API_KEY: $OPENAI_API_KEY

.envにはOpenAIのAPIキーを設定しておきます。

OPENAI_API_KEY=sk-YOUR-API-KEY

Dockerfile
langchainのバージョン変更で頻繁に動かなくなったのでライブラリのバージョンを固定しています。

FROM continuumio/anaconda3

# 仮想環境の構築
RUN conda create -n faiss-env && \
    echo "conda activate faiss-env" >> ~/.bashrc
SHELL ["/bin/bash", "--login", "-c"]

# 必要なパッケージをインストール
RUN conda install -c conda-forge faiss==1.7.2

RUN pip install \
    openai==0.27.6 \
    langchain==0.0.161 \
    fastapi==0.95.1 \
    uvicorn==0.22.0 \
    pydantic==1.10.7 \
    tiktoken==0.4.0 \
    lark==1.1.5

WORKDIR /app

# アプリケーションのコピー
COPY ./app /app

EXPOSE 5001

# FastAPIアプリケーションの起動
CMD ["/bin/bash", "-c", "source activate faiss-env && uvicorn main:app --host 0.0.0.0 --port 5001"]

これでDocker上でAnacondaがインストールされ、その実行環境でFastAPIが起動します。

次にAPIサーバの実装です。
ルーティング周りはフロントエンドからAPIを実行するためにCORSの設定をするくらいです。

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from routes import router

app = FastAPI()

# CORS設定
origins = [
    "http://localhost:3000",
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# routerをアプリケーションに組み込む
app.include_router(router)

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=5001)

routes.py

import json
from fastapi import APIRouter
from fastapi.responses import JSONResponse
from models.agent import CreateAgentRequestData, CreateMessageRequestData, CreateMemoryRequestData
from controllers.agent import AgentController

router = APIRouter()
agent_controller = AgentController()

# ルーティング
@router.get('/health')
async def health():
    return {"status": "OK"}

@router.post('/agents', status_code=201)
async def create_agent(input: CreateAgentRequestData):
    agent = agent_controller.create_agent(input)
    return agent.get_summary()

@router.get('/agents', status_code=200)
async def get_agents():
    agents = agent_controller.get_agents()
    return JSONResponse(content=agents)

@router.get('/agents/{id}', status_code=200)
async def get_agent_detail(id:int):
    agent = agent_controller.get_agent_detail(id)
    return JSONResponse(content=agent)

@router.post('/agents/{id}/memories', status_code=201)
async def create_memory(id:int, data: CreateMemoryRequestData):
    print("Router: Create Memory")
    agent_controller.create_memory(id, data.memory)
    return

@router.post('/agents/{id}/messages', status_code=201)
async def create_message(id:int, data: CreateMessageRequestData):
    res = agent_controller.create_message(id, data.message)
    return JSONResponse(content=res)

次にコントローラーです。
IDやエージェントの名前などはRDBに、記憶部分をChromaDBなどのベクトルDBに格納しておくことで永続化が出来そうな気がしますが、まだそこまで検証できていません…とりあえずローカルで動けば良いかと思ってグローバル変数に格納しています…

from models.agent import CreateAgentRequestData, Agent

id = 0
agents = {}

class AgentController:
    def create_agent(self, data: CreateAgentRequestData):
        global id
        id += 1
        agent = Agent(id, data)
        agents[agent.id] = agent
        return agent
    
    def get_agents(self):
        res = {}
        list = []
        for agent in agents.values():
            dict = {}
            dict["id"] = agent.id
            dict["name"] = agent.agent.name
            dict["age"] = agent.agent.age
            dict["traits"] = agent.agent.traits
            dict["status"] = agent.agent.status
            list.append(dict)
        res["agentList"] = list
        return res
        
    def get_agent_detail(self, id):
        res = {}
        if id in agents:
            agent = agents[id]
            summary = agent.get_summary()
            res["id"] = agent.id
            res["name"] = agent.agent.name
            res["age"] = agent.agent.age
            res["traits"] = agent.agent.traits
            res["status"] = agent.agent.status
            res["summary"] = summary
        return res

    def create_memory(self, id, memory):
        if id in agents:
            agent = agents[id]
            agent.add_memory(memory)
        return

    def create_message(self, id, message):
        res = {}
        reply = ""
        if id in agents:
            agent = agents[id]
            reply = agent.interview_agent(message)
        res["reply"] = reply
        return res

最後にモデル層です。
基本的にはLangChainのページに書かれている実装例を参考にしています。

import os
import math
import faiss
from typing import List
from pydantic import BaseModel
from langchain.chat_models import ChatOpenAI
from langchain.docstore import InMemoryDocstore
from langchain.embeddings import OpenAIEmbeddings
from langchain.retrievers import TimeWeightedVectorStoreRetriever
from langchain.vectorstores import FAISS
# from models.generative_agent import GenerativeAgent
# from models.memory import GenerativeAgentMemory
from langchain.experimental.generative_agents import GenerativeAgent, GenerativeAgentMemory

API_KEY=os.environ["OPENAI_API_KEY"]
LLM = ChatOpenAI(max_tokens=1500, openai_api_key=API_KEY)

class CreateAgentRequestData(BaseModel):
    name: str
    age: int
    traits: str
    status: str
    memory: str

class CreateMemoryRequestData(BaseModel):
    memory: str

class CreateMessageRequestData(BaseModel):
    message: str

class Agent:
    def __init__(self, id, data: CreateAgentRequestData):
        self.id = id
        # 記憶の生成
        memory = GenerativeAgentMemory(
            llm=LLM,
            memory_retriever=create_new_memory_retriever(),
            verbose=True,
            reflection_threshold=8
        )
        # エージェント生成
        self.agent = GenerativeAgent(
                    name=data.name, 
                    age=data.age,
                    traits=data.traits, # You can add more persistent traits here 
                    status=data.status, # When connected to a virtual world, we can have the characters update their status
                    llm=LLM,
                    memory=memory
                )
        # 初期記憶
        self.add_memory(data.memory)
            
    def get_summary(self):
        return self.agent.get_summary(force_refresh=True)

    def add_memory(self, memory):
        self.agent.memory.add_memory(memory)

    def interview_agent(self, message):
        return interview_agent(self.agent, message)

def relevance_score_fn(score: float) -> float:
    """Return a similarity score on a scale [0, 1]."""
    # This will differ depending on a few things:
    # - the distance / similarity metric used by the VectorStore
    # - the scale of your embeddings (OpenAI's are unit norm. Many others are not!)
    # This function converts the euclidean norm of normalized embeddings
    # (0 is most similar, sqrt(2) most dissimilar)
    # to a similarity function (0 to 1)
    return 1.0 - score / math.sqrt(2)

def create_new_memory_retriever():
    """Create a new vector store retriever unique to the agent."""
    # Define your embedding model
    embeddings_model = OpenAIEmbeddings()
    # Initialize the vectorstore as empty
    embedding_size = 1536
    index = faiss.IndexFlatL2(embedding_size)
    vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {}, relevance_score_fn=relevance_score_fn)
    return TimeWeightedVectorStoreRetriever(vectorstore=vectorstore, other_score_keys=["importance"], k=15)    

def interview_agent(agent: GenerativeAgent, message: str) -> str:
    """Help the notebook user interact with the agent."""
    return agent.generate_dialogue_response(message)[1]

これでエージェントの作成、記憶の設定、チャットのAPIの実装ができました。 GenerativeAgent, GenerativeAgentMemoryのコードも公開されているので、それを修正すれば独自実装することもできます。

langchain.experimental.generative_agents.generative_agent — 🦜🔗 LangChain 0.0.165

langchain.experimental.generative_agents.memory — 🦜🔗 LangChain 0.0.165

以上でAPIの実装が出来ました。
あとはクライアントをよしなに実装してAPIを叩いてやれば完成です。

おわりに

以上、GenerativeAgentsを操作するAPIを実装してみました。
記憶の永続化なども実装していきたいのですが、一旦ローカルで手軽に遊べるようになったのでここまでにします。

このブログでは他にもAI関連の記事もありますのでぜひご覧ください!

synamon.hatenablog.com

synamon.hatenablog.com

synamon.hatenablog.com

synamon.hatenablog.com