日経Linux2023年11月号LLM特集の生成エージェントリポジトリ
from __future__ import annotations

from json import JSONDecodeError
from typing import Any

from langchain.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)
from langchain.chains import LLMChain
from langchain.chains.base import Chain
from langchain.docstore.document import Document
from langchain.output_parsers import RetryWithErrorOutputParser
from langchain.prompts import PromptTemplate
from langchain.prompts.base import BasePromptTemplate
from langchain.pydantic_v1 import Field
from langchain.schema.language_model import BaseLanguageModel
from pydantic import Extra

from retriever import TimeImportanceSimilarityRetriever


class GAChain(Chain):
    prompt: BasePromptTemplate
    observation_prompt: BasePromptTemplate
    importance_chain: LLMChain
    llm: BaseLanguageModel
    retriever: TimeImportanceSimilarityRetriever
    default_importance = 3.0
    parser: RetryWithErrorOutputParser

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid
        arbitrary_types_allowed = True

    @property
    def input_keys(self) -> list[str]:
        """Will be whatever keys the prompt expects.

        :meta private:
        """
        return ["input_text", "speaker"]

    @property
    def output_keys(self) -> list[str]:
        """Will always return text key.

        :meta private:
        """
        return ["text", "input_importance", "output_importance"]

    def memorize(self, observation: str, context: str) -> tuple[Document, float]:
        result = self.importance_chain.run(
            context=context,
            observation=observation,
        )
        prompt_value = self.importance_chain.prompt.format_prompt(
            context=context,
            observation=observation,
        )
        try:
            importance = self.parser.parse_with_prompt(result, prompt_value).rate
        except JSONDecodeError:
            importance = self.default_importance

        doc = Document(page_content=observation, metadata={"importance": importance})
        self.retriever.add_documents([doc])
        return doc, importance

    def get_context(self, query: str) -> str:
        docs = self.retriever.get_relevant_documents(query)
        docs = sorted(docs, key=lambda x: x.metadata["created_time"])
        context = "\n".join([doc.page_content for doc in docs])
        return context

    def _call(
        self,
        inputs: dict[str, Any],
        run_manager: CallbackManagerForChainRun | None = None,
    ) -> dict[str, str]:
        observation = self.observation_prompt.format_prompt(**inputs).text
        context = self.get_context(observation)
        _, input_importance = self.memorize(observation, context)
        prompt_value = self.prompt.format_prompt(
            context=context,
            observation=observation,
            name=self.retriever.agent_name,
        )
        response = self.llm.generate_prompt(
            [prompt_value], callbacks=run_manager.get_child() if run_manager else None
        )

        context += "\n" + observation
        generated_text = response.generations[0][0].text
        observation = self.observation_prompt.format_prompt(
            speaker=self.retriever.agent_name, input_text=generated_text
        ).text
        _, output_importance = self.memorize(observation, context)

        if run_manager:
            run_manager.on_text("Log something about this run")

        return {
            "text": generated_text,
            "input_importance": input_importance,
            "output_importance": output_importance,
        }

    @property
    def _chain_type(self) -> str:
        return "Generative_Agents_Chain"

chain.py
elith-co-jp/nikkei-linux-2023-11-generative-agents
import gradio as gr
import pandas as pd
from dotenv import load_dotenv
from faiss import IndexFlatIP
from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.docstore import InMemoryDocstore
from langchain.docstore.document import Document
from langchain.embeddings import OpenAIEmbeddings
from langchain.output_parsers import PydanticOutputParser, RetryWithErrorOutputParser
from langchain.prompts import PromptTemplate
from langchain.vectorstores import FAISS
from pydantic import BaseModel, Field

from chain import GAChain
from retriever import TimeImportanceSimilarityRetriever

load_dotenv()

IMPORTANCE_TEMPLATE = """\
あなたの名前は Elith bot です。
以下の Observation に示す内容の発言がありました。
この発言に関連する内容を Context に列挙します。
Observation の発言の重要度について 1 から 10 で客観的に評価してください。
目安として、1 は歯を磨く、ベッドを整えるなどの日常的な行為、10 は別れ、大学合格などの非常に重要な出来事とします。

Context: 
```{context}```
Observation: 
```{observation}```

出力は以下のJSONスキーマに準拠したJSONインスタンスとしてフォーマットしてください。

例として、スキーマ{{"properties": {{"foo": {{"title": "Foo", "description": "a list of strings", "type": "array", "items": {{"type": "string"}}}}}}, "required": ["foo"]}}に対して、オブジェクト{{"foo": ["bar", "baz"]}}はスキーマの適切にフォーマットされたインスタンスです。しかし、オブジェクト{{"properties": {{"foo": ["bar", "baz"]}}}}は適切にフォーマットされていません。

今回用いるスキーマは以下の通りです。
{{"properties": {{"rate": {{"description": "重要度を表す 1 から 10 の整数", "title": "Rate", "type": "integer"}}}}, "required": ["rate"]}}

出力:
"""

OBSERVATION_TEMPLATE = "{speaker}: {input_text}"

TEMPLATE = """\
あなたの名前は {name} です。

Context として、関連する発言を示します。
Context: 
```{context}```

Context を考慮して以下の発言に対して、返答してください。
```{observation}```

{name}:
"""


def create_demo(chain: GAChain):
    def respond(speaker, message, chat_history, history):
        result = chain(dict(input_text=message, speaker=speaker))
        bot_message = result["text"]
        input_importance = result["input_importance"]
        output_importance = result["output_importance"]
        chat_history.append((message, bot_message))
        history["name"].append(speaker)
        history["importance"].append(input_importance)
        history["text"].append(message)
        history["name"].append(chain.retriever.agent_name)
        history["importance"].append(output_importance)
        history["text"].append(bot_message)

        return "", chat_history, history

    def create_df(history):
        df = pd.DataFrame(history)
        return df

    with gr.Blocks() as demo:
        name = gr.Text(label="あなたの名前", interactive=True)
        chatbot = gr.Chatbot()
        msg = gr.Textbox()
        clear = gr.ClearButton([msg, chatbot])
        history = gr.State(value={"name": [], "importance": [], "text": []})

        msg.submit(respond, [name, msg, chatbot, history], [msg, chatbot, history])
        with gr.Accordion(label="チャット履歴", open=False):
            button = gr.Button("履歴を取得")
            df = gr.DataFrame()
        button.click(create_df, inputs=history, outputs=df)

    demo.launch()


def main():
    llm_0 = ChatOpenAI(temperature=0)
    llm = ChatOpenAI()

    embeddings_model = OpenAIEmbeddings()
    embedding_size = 1536
    index = IndexFlatIP(embedding_size)
    vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
    retriever = TimeImportanceSimilarityRetriever(
        vectorstore=vectorstore,
        decay_rate=0.005,
        k=5,
        weights=dict(importance=0.1, similarity=1, time=1),
        agent_name="Elith bot",
    )

    importance_prompt = PromptTemplate(
        template=IMPORTANCE_TEMPLATE, input_variables=["context", "observation"]
    )
    importance_chain = LLMChain(llm=llm_0, prompt=importance_prompt)
    observation_prompt = PromptTemplate(
        template=OBSERVATION_TEMPLATE, input_variables=["speaker", "input_text"]
    )
    prompt = PromptTemplate(
        template=TEMPLATE, input_variables=["name", "context", "observation"]
    )

    class ImportanceRate(BaseModel):
        rate: int = Field(description="重要度を表す 1 から 10 の整数")

    parser = PydanticOutputParser(pydantic_object=ImportanceRate)
    retry_parser = RetryWithErrorOutputParser.from_llm(parser=parser, llm=llm_0)

    chain = GAChain(
        prompt=prompt,
        observation_prompt=observation_prompt,
        retriever=retriever,
        llm=llm,
        importance_chain=importance_chain,
        parser=retry_parser,
    )

    create_demo(chain)


if __name__ == "__main__":
    main()

main.py
elith-co-jp/nikkei-linux-2023-11-generative-agents
from __future__ import annotations

import datetime
from copy import deepcopy
from typing import Any, TypeVar

from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import Field
from langchain.schema import BaseRetriever, Document
from langchain.vectorstores.base import VectorStore

Time = TypeVar("Time", int, float, datetime.datetime)


class TimeImportanceSimilarityRetriever(BaseRetriever):
    vectorstore: VectorStore
    agent_name: str
    search_kwargs: dict = Field(default_factory=lambda: dict(k=100))
    memory_stream: list[Document] = Field(default_factory=list)
    decay_rate: float = Field(default=0.01)
    weights: dict[str, float] = Field(
        default_factory=lambda: dict(similarity=1.0, time=1.0, importance=1.0)
    )
    k: int = 4
    use_datetime: bool = False
    current_time: int = 0

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:
        if self.use_datetime:
            current_time = datetime.datetime.now()
        else:
            current_time = self.current_time

        docs_scores = self._get_combined_score(query, current_time=current_time)
        sorted_docs_scores = sorted(docs_scores, key=lambda x: x[1], reverse=True)
        documents = [doc for doc, _ in sorted_docs_scores[: self.k]]

        # update last_access_time
        for document in documents:
            document.metadata["last_access_time"] = current_time
        return documents

    def _get_combined_score(
        self, query: str, current_time: Time
    ) -> list[tuple[Document, float]]:
        """時間・重要度・類似度を合算したスコアを計算する

        Args:
            query (str): クエリテキスト
            current_time (Time): 現在時刻

        Returns:
            list(tuple[Document, float]): ドキュメントとスコアのリスト
        """
        docs_and_scores = self.vectorstore.similarity_search_with_relevance_scores(
            query, **self.search_kwargs
        )
        docs_total_scores = []
        for doc, similarity in docs_and_scores:
            time_passed = self._get_time_passed(
                current_time, doc.metadata["last_access_time"]
            )
            score_time = (1.0 - self.decay_rate) ** time_passed
            score_importance = doc.metadata["importance"]
            total_score = (
                self.weights["similarity"] * similarity
                + self.weights["time"] * score_time
                + self.weights["importance"] * score_importance
            )
            docs_total_scores.append((doc, total_score))
        return docs_total_scores

    def _get_time_passed(self, current_time: Time, last_access_time: Time) -> Time:
        """時間の差を計算する

        Args:
            current_time (Time): 現在時刻
            last_access_time (Time): 最終アクセス時刻

        Returns:
            Time: 時間の差
        """
        if isinstance(current_time, datetime.datetime):
            hours_passed = (current_time - last_access_time).total_seconds() / 3600
            return hours_passed
        return current_time - last_access_time

    def add_documents(self, documents: list[Document], **kwargs: Any) -> list[str]:
        """Add documents to vectorstore."""

        if self.use_datetime:
            current_time = datetime.datetime.now()
        else:
            current_time = self.current_time

        docs = deepcopy(documents)
        for doc in docs:
            doc.metadata["last_access_time"] = current_time
            doc.metadata["created_time"] = current_time
            if "importance" not in doc.metadata:
                doc.metadata["importance"] = 0
            self.current_time += 1

        return self.vectorstore.add_documents(docs, **kwargs)

retriever.py
elith-co-jp/nikkei-linux-2023-11-generative-agents