Files
langchain/libs/community/langchain_community/tools/memorize/tool.py
Harrison Chase 8516a03a02 langchain-community[major]: Upgrade community to pydantic 2 (#26011)
This PR upgrades langchain-community to pydantic 2.


* Most of this PR was auto-generated using code mods with gritql
(https://github.com/eyurtsev/migrate-pydantic/tree/main)
* Subsequently, some code was fixed manually due to accommodate
differences between pydantic 1 and 2

Breaking Changes:

- Use TEXTEMBED_API_KEY and TEXTEMBEB_API_URL for env variables for text
embed integrations:
cbea780492

Other changes:

- Added pydantic_settings as a required dependency for community. This
may be removed if we have enough time to convert the dependency into an
optional one.

---------

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
Co-authored-by: Bagatur <baskaryan@gmail.com>
2024-09-05 14:07:10 -04:00

61 lines
1.8 KiB
Python

from abc import abstractmethod
from typing import Any, Optional, Protocol, Sequence, runtime_checkable
from langchain_core.callbacks import (
AsyncCallbackManagerForToolRun,
CallbackManagerForToolRun,
)
from langchain_core.tools import BaseTool
from pydantic import Field
from langchain_community.llms.gradient_ai import TrainResult
@runtime_checkable
class TrainableLLM(Protocol):
"""Protocol for trainable language models."""
@abstractmethod
def train_unsupervised(
self,
inputs: Sequence[str],
**kwargs: Any,
) -> TrainResult: ...
@abstractmethod
async def atrain_unsupervised(
self,
inputs: Sequence[str],
**kwargs: Any,
) -> TrainResult: ...
class Memorize(BaseTool):
"""Tool that trains a language model."""
name: str = "memorize"
description: str = (
"Useful whenever you observed novel information "
"from previous conversation history, "
"i.e., another tool's action outputs or human comments. "
"The action input should include observed information in detail, "
"then the tool will fine-tune yourself to remember it."
)
llm: TrainableLLM = Field()
def _run(
self,
information_to_learn: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> str:
train_result = self.llm.train_unsupervised((information_to_learn,))
return f"Train complete. Loss: {train_result['loss']}"
async def _arun(
self,
information_to_learn: str,
run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
) -> str:
train_result = await self.llm.atrain_unsupervised((information_to_learn,))
return f"Train complete. Loss: {train_result['loss']}"