Compare commits
36 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
abdfb1abbf | ||
|
|
2f07248211 | ||
|
|
9ae6fc822b | ||
|
|
5094359c4e | ||
|
|
28801b54bf | ||
|
|
4d69f0e837 | ||
|
|
d91b3f8042 | ||
|
|
03a0ab2fcf | ||
|
|
d860d62972 | ||
|
|
add4893939 | ||
|
|
eb251d6e37 | ||
|
|
4ba635497b | ||
|
|
bdbf14c1ed | ||
|
|
c340fc085c | ||
|
|
94f8d0f242 | ||
|
|
1d8a9b113c | ||
|
|
1837796852 | ||
|
|
c5c53c056e | ||
|
|
f1b1190f0b | ||
|
|
303ec2a139 | ||
|
|
0303a42699 | ||
|
|
d768713630 | ||
|
|
33b043b56d | ||
|
|
b6f4069e8c | ||
|
|
1cb5869f0b | ||
|
|
8ad794aa6c | ||
|
|
611f77a730 | ||
|
|
0738899ac9 | ||
|
|
76b7b5c0e0 | ||
|
|
cae5e8d217 | ||
|
|
97bfd3a672 | ||
|
|
8fb1341b93 | ||
|
|
cba4412f3d | ||
|
|
2ad87f6e8a | ||
|
|
496eb469db | ||
|
|
b78fec3bac |
9
.gitignore
vendored
@@ -7,7 +7,12 @@
|
||||
*_internal*
|
||||
talemate_env
|
||||
chroma
|
||||
scenes
|
||||
config.yaml
|
||||
!scenes/infinity-quest/assets
|
||||
templates/llm-prompt/user/*.jinja2
|
||||
scenes/
|
||||
!scenes/infinity-quest-dynamic-scenario/
|
||||
!scenes/infinity-quest-dynamic-scenario/assets/
|
||||
!scenes/infinity-quest-dynamic-scenario/templates/
|
||||
!scenes/infinity-quest-dynamic-scenario/infinity-quest.json
|
||||
!scenes/infinity-quest/assets/
|
||||
!scenes/infinity-quest/infinity-quest.json
|
||||
|
||||
208
README.md
@@ -1,36 +1,56 @@
|
||||
# Talemate
|
||||
|
||||
Allows you to play roleplay scenarios with large language models.
|
||||
Roleplay with AI with a focus on strong narration and consistent world and game state tracking.
|
||||
|
||||
It does not run any large language models itself but relies on existing APIs. Currently supports **text-generation-webui** and **openai**.
|
||||
|||
|
||||
|------------------------------------------|------------------------------------------|
|
||||
|||
|
||||
|||
|
||||
|
||||
This means you need to either have an openai api key or know how to setup [oobabooga/text-generation-webui](https://github.com/oobabooga/text-generation-webui) (locally or remotely via gpu renting. `--extension openai` flag needs to be set)
|
||||
> :warning: **It does not run any large language models itself but relies on existing APIs. Currently supports OpenAI, Anthropic, mistral.ai, self-hosted text-generation-webui and LMStudio. 0.18.0 also adds support for generic OpenAI api implementations, but generation quality on that will vary.**
|
||||
|
||||
As of version 0.13.0 the legacy text-generator-webui API `--extension api` is no longer supported, please use their new `--extension openai` api implementation instead.
|
||||
Officially supported APIs:
|
||||
- [OpenAI](https://platform.openai.com/overview)
|
||||
- [Anthropic](https://www.anthropic.com/)
|
||||
- [mistral.ai](https://mistral.ai/)
|
||||
|
||||

|
||||

|
||||
Officially supported self-hosted APIs:
|
||||
- [oobabooga/text-generation-webui](https://github.com/oobabooga/text-generation-webui) (local or with runpod support)
|
||||
- [LMStudio](https://lmstudio.ai/)
|
||||
|
||||
Generic OpenAI api implementations (tested and confirmed working):
|
||||
- [DeepInfra](https://deepinfra.com/) - see [instructions](https://github.com/vegu-ai/talemate/issues/78#issuecomment-1986884304)
|
||||
- [llamacpp](https://github.com/ggerganov/llama.cpp) with the `api_like_OAI.py` wrapper
|
||||
- let me know if you have tested any other implementations and they failed / worked or landed somewhere in between
|
||||
|
||||
## Current features
|
||||
|
||||
- responive modern ui
|
||||
- responsive modern ui
|
||||
- agents
|
||||
- conversation: handles character dialogue
|
||||
- narration: handles narrative exposition
|
||||
- summarization: handles summarization to compress context while maintain history
|
||||
- summarization: handles summarization to compress context while maintaining history
|
||||
- director: can be used to direct the story / characters
|
||||
- editor: improves AI responses (very hit and miss at the moment)
|
||||
- world state: generates world snapshot and handles passage of time (objects and characters)
|
||||
- creator: character / scenario creator
|
||||
- multi-client (agents can be connected to separate APIs)
|
||||
- tts: text to speech via elevenlabs, OpenAI or local tts
|
||||
- visual: stable-diffusion client for in place visual generation via AUTOMATIC1111, ComfyUI or OpenAI
|
||||
- multi-client support (agents can be connected to separate APIs)
|
||||
- long term memory
|
||||
- chromadb integration
|
||||
- passage of time
|
||||
- narrative world state
|
||||
- Automatically keep track and reinforce selected character and world truths / states.
|
||||
- narrative tools
|
||||
- creative tools
|
||||
- manage multiple NPCs
|
||||
- AI backed character creation with template support (jinja2)
|
||||
- AI backed scenario creation
|
||||
- context managegement
|
||||
- Manage character details and attributes
|
||||
- Manage world information / past events
|
||||
- Pin important information to the context (Manually or conditionally through AI)
|
||||
- runpod integration
|
||||
- overridable templates for all prompts. (jinja2)
|
||||
|
||||
@@ -46,98 +66,159 @@ In no particular order:
|
||||
- Dynamic player choice generation
|
||||
- Better creative tools
|
||||
- node based scenario / character creation
|
||||
- Improved and consistent long term memory
|
||||
- Improved and consistent long term memory and accurate current state of the world
|
||||
- Improved director agent
|
||||
- Right now this doesn't really work well on anything but GPT-4 (and even there it's debatable). It tends to steer the story in a way that introduces pacing issues. It needs a model that is creative but also reasons really well i think.
|
||||
- Gameplay loop governed by AI
|
||||
- objectives
|
||||
- quests
|
||||
- win / lose conditions
|
||||
- Automatic1111 client for in place visual generation
|
||||
|
||||
|
||||
# Instructions
|
||||
|
||||
Please read the documents in the `docs` folder for more advanced configuration and usage.
|
||||
|
||||
- [Quickstart](#quickstart)
|
||||
- [Installation](#installation)
|
||||
- [Connecting to an LLM](#connecting-to-an-llm)
|
||||
- [Text-generation-webui](#text-generation-webui)
|
||||
- [Recommended Models](#recommended-models)
|
||||
- [OpenAI / mistral.ai / Anthropic](#openai)
|
||||
- [DeepInfra via OpenAI Compatible client](#deepinfra-via-openai-compatible-client)
|
||||
- [Ready to go](#ready-to-go)
|
||||
- [Load the introductory scenario "Infinity Quest"](#load-the-introductory-scenario-infinity-quest)
|
||||
- [Loading character cards](#loading-character-cards)
|
||||
- [Text-to-Speech (TTS)](docs/tts.md)
|
||||
- [Visual Generation](docs/visual.md)
|
||||
- [ChromaDB (long term memory) configuration](docs/chromadb.md)
|
||||
- [Runpod Integration](docs/runpod.md)
|
||||
- [Prompt template overrides](docs/templates.md)
|
||||
|
||||
# Quickstart
|
||||
|
||||
## Installation
|
||||
|
||||
Post [here](https://github.com/final-wombat/talemate/issues/17) if you run into problems during installation.
|
||||
Post [here](https://github.com/vegu-ai/talemate/issues/17) if you run into problems during installation.
|
||||
|
||||
There is also a [troubleshooting guide](docs/troubleshoot.md) that might help.
|
||||
|
||||
### Windows
|
||||
|
||||
1. Download and install Python 3.10 or higher from the [official Python website](https://www.python.org/downloads/windows/).
|
||||
1. Download and install Node.js from the [official Node.js website](https://nodejs.org/en/download/). This will also install npm.
|
||||
1. Download the Talemate project to your local machine. Download from [the Releases page](https://github.com/final-wombat/talemate/releases).
|
||||
1. Download and install Python 3.10 or Python 3.11 from the [official Python website](https://www.python.org/downloads/windows/). :warning: python3.12 is currently not supported.
|
||||
1. Download and install Node.js v20 from the [official Node.js website](https://nodejs.org/en/download/). This will also install npm. :warning: v21 is currently not supported.
|
||||
1. Download the Talemate project to your local machine. Download from [the Releases page](https://github.com/vegu-ai/talemate/releases).
|
||||
1. Unpack the download and run `install.bat` by double clicking it. This will set up the project on your local machine.
|
||||
1. Once the installation is complete, you can start the backend and frontend servers by running `start.bat`.
|
||||
1. Navigate your browser to http://localhost:8080
|
||||
|
||||
### Linux
|
||||
|
||||
`python 3.10` or higher is required.
|
||||
`python 3.10` or `python 3.11` is required. :warning: `python 3.12` not supported yet.
|
||||
|
||||
1. `git clone git@github.com:final-wombat/talemate`
|
||||
`nodejs v19 or v20` :warning: `v21` not supported yet.
|
||||
|
||||
1. `git clone git@github.com:vegu-ai/talemate`
|
||||
1. `cd talemate`
|
||||
1. `source install.sh`
|
||||
1. Start the backend: `python src/talemate/server/run.py runserver --host 0.0.0.0 --port 5050`.
|
||||
1. Open a new terminal, navigate to the `talemate_frontend` directory, and start the frontend server by running `npm run serve`.
|
||||
|
||||
## Configuration
|
||||
|
||||
### OpenAI
|
||||
|
||||
To set your openai api key, open `config.yaml` in any text editor and uncomment / add
|
||||
|
||||
```yaml
|
||||
openai:
|
||||
api_key: sk-my-api-key-goes-here
|
||||
```
|
||||
|
||||
You will need to restart the backend for this change to take effect.
|
||||
|
||||
### RunPod
|
||||
|
||||
To set your runpod api key, open `config.yaml` in any text editor and uncomment / add
|
||||
|
||||
```yaml
|
||||
runpod:
|
||||
api_key: my-api-key-goes-here
|
||||
```
|
||||
You will need to restart the backend for this change to take effect.
|
||||
|
||||
Once the api key is set Pods loaded from text-generation-webui templates (or the bloke's runpod llm template) will be autoamtically added to your client list in talemate.
|
||||
|
||||
**ATTENTION**: Talemate is not a suitable for way for you to determine whether your pod is currently running or not. **Always** check the runpod dashboard to see if your pod is running or not.
|
||||
|
||||
## Recommended Models
|
||||
(as of2023.10.25)
|
||||
|
||||
Any of the top models in any of the size classes here should work well:
|
||||
https://www.reddit.com/r/LocalLLaMA/comments/17fhp9k/huge_llm_comparisontest_39_models_tested_7b70b/
|
||||
|
||||
## Connecting to an LLM
|
||||
# Connecting to an LLM
|
||||
|
||||
On the right hand side click the "Add Client" button. If there is no button, you may need to toggle the client options by clicking this button:
|
||||
|
||||
As of version 0.13.0 the legacy text-generator-webui API `--extension api` is no longer supported, please use their new `--extension openai` api implementation instead.
|
||||
|
||||

|
||||
|
||||
### Text-generation-webui
|
||||

|
||||
|
||||
## Text-generation-webui
|
||||
|
||||
> :warning: As of version 0.13.0 the legacy text-generator-webui API `--extension api` is no longer supported, please use their new `--extension openai` api implementation instead.
|
||||
|
||||
In the modal if you're planning to connect to text-generation-webui, you can likely leave everything as is and just click Save.
|
||||
|
||||

|
||||

|
||||
|
||||
### OpenAI
|
||||
### Specifying the correct prompt template
|
||||
|
||||
For good results it is **vital** that the correct prompt template is specified for whichever model you have loaded.
|
||||
|
||||
Talemate does come with a set of pre-defined templates for some popular models, but going forward, due to the sheet number of models released every day, understanding and specifying the correct prompt template is something you should familiarize yourself with.
|
||||
|
||||
If the text-gen-webui client shows a yellow triangle next to it, it means that the prompt template is not set, and it is currently using the default `VICUNA` style prompt template.
|
||||
|
||||

|
||||
|
||||
Click the two cogwheels to the right of the triangle to open the client settings.
|
||||
|
||||

|
||||
|
||||
You can first try by clicking the `DETERMINE VIA HUGGINGFACE` button, depending on the model's README file, it may be able to determine the correct prompt template for you. (basically the readme needs to contain an example of the template)
|
||||
|
||||
If that doesn't work, you can manually select the prompt template from the dropdown.
|
||||
|
||||
In the case for `bartowski_Nous-Hermes-2-Mistral-7B-DPO-exl2_8_0` that is `ChatML` - select it from the dropdown and click `Save`.
|
||||
|
||||

|
||||
|
||||
### Recommended Models
|
||||
|
||||
As of 2024.03.07 my personal regular drivers (the ones i test with) are:
|
||||
|
||||
- Kunoichi-7B
|
||||
- sparsetral-16x7B
|
||||
- Nous-Hermes-2-Mistral-7B-DPO
|
||||
- brucethemoose_Yi-34B-200K-RPMerge
|
||||
- dolphin-2.7-mixtral-8x7b
|
||||
- rAIfle_Verdict-8x7B
|
||||
- Mixtral-8x7B-instruct
|
||||
|
||||
That said, any of the top models in any of the size classes here should work well (i wouldn't recommend going lower than 7B):
|
||||
|
||||
https://www.reddit.com/r/LocalLLaMA/comments/18yp9u4/llm_comparisontest_api_edition_gpt4_vs_gemini_vs/
|
||||
|
||||
## OpenAI / mistral.ai / Anthropic
|
||||
|
||||
The setup is the same for all three, the example below is for OpenAI.
|
||||
|
||||
If you want to add an OpenAI client, just change the client type and select the apropriate model.
|
||||
|
||||

|
||||

|
||||
|
||||
### Ready to go
|
||||
If you are setting this up for the first time, you should now see the client, but it will have a red dot next to it, stating that it requires an API key.
|
||||
|
||||

|
||||
|
||||
Click the `SET API KEY` button. This will open a modal where you can enter your API key.
|
||||
|
||||

|
||||
|
||||
Click `Save` and after a moment the client should have a green dot next to it, indicating that it is ready to go.
|
||||
|
||||

|
||||
|
||||
## DeepInfra via OpenAI Compatible client
|
||||
|
||||
You can use the OpenAI compatible client to connect to [DeepInfra](https://deepinfra.com/).
|
||||
|
||||

|
||||
|
||||
```
|
||||
API URL: https://api.deepinfra.com/v1/openai
|
||||
```
|
||||
|
||||
Models on DeepInfra that work well with Talemate:
|
||||
|
||||
- [mistralai/Mixtral-8x7B-Instruct-v0.1](https://deepinfra.com/mistralai/Mixtral-8x7B-Instruct-v0.1) (max context 32k, 8k recommended)
|
||||
- [cognitivecomputations/dolphin-2.6-mixtral-8x7b](https://deepinfra.com/cognitivecomputations/dolphin-2.6-mixtral-8x7b) (max context 32k, 8k recommended)
|
||||
- [lizpreciatior/lzlv_70b_fp16_hf](https://deepinfra.com/lizpreciatior/lzlv_70b_fp16_hf) (max context 4k)
|
||||
|
||||
## Ready to go
|
||||
|
||||
You will know you are good to go when the client and all the agents have a green dot next to them.
|
||||
|
||||

|
||||

|
||||
|
||||
## Load the introductory scenario "Infinity Quest"
|
||||
|
||||
@@ -157,11 +238,4 @@ Expand the "Load" menu in the top left corner and either click on "Upload a char
|
||||
|
||||
Once a character is uploaded, talemate may actually take a moment because it needs to convert it to a talemate format and will also run additional LLM prompts to generate character attributes and world state.
|
||||
|
||||
Make sure you save the scene after the character is loaded as it can then be loaded as normal talemate scenario in the future.
|
||||
|
||||
## Further documentation
|
||||
|
||||
- Creative mode (docs WIP)
|
||||
- Prompt template overrides
|
||||
- [ChromaDB (long term memory)](docs/chromadb.md)
|
||||
- Runpod Integration
|
||||
Make sure you save the scene after the character is loaded as it can then be loaded as normal talemate scenario in the future.
|
||||
@@ -2,25 +2,73 @@ agents: {}
|
||||
clients: {}
|
||||
creator:
|
||||
content_context:
|
||||
- a fun and engaging slice of life story aimed at an adult audience.
|
||||
- a terrifying horror story aimed at an adult audience.
|
||||
- a thrilling action story aimed at an adult audience.
|
||||
- a mysterious adventure aimed at an adult audience.
|
||||
- an epic sci-fi adventure aimed at an adult audience.
|
||||
- a fun and engaging slice of life story
|
||||
- a terrifying horror story
|
||||
- a thrilling action story
|
||||
- a mysterious adventure
|
||||
- an epic sci-fi adventure
|
||||
game:
|
||||
default_player_character:
|
||||
color: '#6495ed'
|
||||
description: a young man with a penchant for adventure.
|
||||
gender: male
|
||||
name: Elmer
|
||||
world_state:
|
||||
templates:
|
||||
state_reinforcement:
|
||||
Goals:
|
||||
auto_create: false
|
||||
description: Long term and short term goals
|
||||
favorite: true
|
||||
insert: conversation-context
|
||||
instructions: Create a long term goal and two short term goals for {character_name}. Your response must only be the long terms and two short term goals.
|
||||
interval: 20
|
||||
name: Goals
|
||||
query: Goals
|
||||
state_type: npc
|
||||
Physical Health:
|
||||
auto_create: false
|
||||
description: Keep track of health.
|
||||
favorite: true
|
||||
insert: sequential
|
||||
instructions: ''
|
||||
interval: 10
|
||||
name: Physical Health
|
||||
query: What is {character_name}'s current physical health status?
|
||||
state_type: character
|
||||
Time of day:
|
||||
auto_create: false
|
||||
description: Track night / day cycle
|
||||
favorite: true
|
||||
insert: sequential
|
||||
instructions: ''
|
||||
interval: 10
|
||||
name: Time of day
|
||||
query: What is the current time of day?
|
||||
state_type: world
|
||||
|
||||
## Long-term memory
|
||||
|
||||
#chromadb:
|
||||
# embeddings: instructor
|
||||
# instructor_device: cuda
|
||||
# instructor_model: hkunlp/instructor-xl
|
||||
# openai_model: text-embedding-3-small
|
||||
|
||||
## Remote LLMs
|
||||
|
||||
#openai:
|
||||
# api_key: <API_KEY>
|
||||
|
||||
#runpod:
|
||||
# api_key: <API_KEY>
|
||||
# api_key: <API_KEY>
|
||||
|
||||
## TTS (Text-to-Speech)
|
||||
|
||||
#elevenlabs:
|
||||
# api_key: <API_KEY>
|
||||
|
||||
#coqui:
|
||||
# api_key: <API_KEY>
|
||||
|
||||
#tts:
|
||||
# device: cuda
|
||||
# model: tts_models/multilingual/multi-dataset/xtts_v2
|
||||
# voices:
|
||||
# - label: <name>
|
||||
# value: <path to .wav for voice sample>
|
||||
@@ -56,6 +56,7 @@ Then add the following to `config.yaml` for chromadb:
|
||||
```yaml
|
||||
chromadb:
|
||||
embeddings: openai
|
||||
openai_model: text-embedding-3-small
|
||||
```
|
||||
|
||||
**Note**: As with everything openai, using this isn't free. It's way cheaper than their text completion though. ALSO - if you send super explicit content they may flag / ban your key, so keep that in mind (i hear they usually send warnings first though), and always monitor your usage on their dashboard.
|
||||
48
docs/dev/agents/example/test/__init__.py
Normal file
@@ -0,0 +1,48 @@
|
||||
from talemate.agents.base import Agent, AgentAction
|
||||
from talemate.agents.registry import register
|
||||
from talemate.events import GameLoopEvent
|
||||
import talemate.emit.async_signals
|
||||
from talemate.emit import emit
|
||||
|
||||
@register()
|
||||
class TestAgent(Agent):
|
||||
|
||||
agent_type = "test"
|
||||
verbose_name = "Test"
|
||||
|
||||
def __init__(self, client):
|
||||
self.client = client
|
||||
self.is_enabled = True
|
||||
self.actions = {
|
||||
"test": AgentAction(
|
||||
enabled=True,
|
||||
label="Test",
|
||||
description="Test",
|
||||
),
|
||||
}
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return self.is_enabled
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
return True
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return True
|
||||
|
||||
def connect(self, scene):
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("game_loop").connect(self.on_game_loop)
|
||||
|
||||
async def on_game_loop(self, emission: GameLoopEvent):
|
||||
"""
|
||||
Called on the beginning of every game loop
|
||||
"""
|
||||
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
emit("status", status="info", message="Annoying you with a test message every game loop.")
|
||||
130
docs/dev/client/example/runpod_vllm/__init__.py
Normal file
@@ -0,0 +1,130 @@
|
||||
"""
|
||||
An attempt to write a client against the runpod serverless vllm worker.
|
||||
|
||||
This is close to functional, but since runpod serverless gpu availability is currently terrible, i have
|
||||
been unable to properly test it.
|
||||
|
||||
Putting it here for now since i think it makes a decent example of how to write a client against a new service.
|
||||
"""
|
||||
|
||||
import pydantic
|
||||
import structlog
|
||||
import runpod
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from talemate.client.base import ClientBase, ExtraField
|
||||
from talemate.client.registry import register
|
||||
from talemate.emit import emit
|
||||
from talemate.config import Client as BaseClientConfig
|
||||
|
||||
log = structlog.get_logger("talemate.client.runpod_vllm")
|
||||
|
||||
class Defaults(pydantic.BaseModel):
|
||||
max_token_length: int = 4096
|
||||
model: str = ""
|
||||
runpod_id: str = ""
|
||||
|
||||
class ClientConfig(BaseClientConfig):
|
||||
runpod_id: str = ""
|
||||
|
||||
@register()
|
||||
class RunPodVLLMClient(ClientBase):
|
||||
client_type = "runpod_vllm"
|
||||
conversation_retries = 5
|
||||
config_cls = ClientConfig
|
||||
|
||||
class Meta(ClientBase.Meta):
|
||||
title: str = "Runpod VLLM"
|
||||
name_prefix: str = "Runpod VLLM"
|
||||
enable_api_auth: bool = True
|
||||
manual_model: bool = True
|
||||
defaults: Defaults = Defaults()
|
||||
extra_fields: dict[str, ExtraField] = {
|
||||
"runpod_id": ExtraField(
|
||||
name="runpod_id",
|
||||
type="text",
|
||||
label="Runpod ID",
|
||||
required=True,
|
||||
description="The Runpod ID to connect to.",
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
def __init__(self, model=None, runpod_id=None, **kwargs):
|
||||
self.model_name = model
|
||||
self.runpod_id = runpod_id
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return False
|
||||
|
||||
|
||||
def set_client(self, **kwargs):
|
||||
log.debug("set_client", kwargs=kwargs, runpod_id=self.runpod_id)
|
||||
self.runpod_id = kwargs.get("runpod_id", self.runpod_id)
|
||||
|
||||
|
||||
def tune_prompt_parameters(self, parameters: dict, kind: str):
|
||||
super().tune_prompt_parameters(parameters, kind)
|
||||
|
||||
keys = list(parameters.keys())
|
||||
|
||||
valid_keys = ["temperature", "top_p", "max_tokens"]
|
||||
|
||||
for key in keys:
|
||||
if key not in valid_keys:
|
||||
del parameters[key]
|
||||
|
||||
async def get_model_name(self):
|
||||
return self.model_name
|
||||
|
||||
async def generate(self, prompt: str, parameters: dict, kind: str):
|
||||
"""
|
||||
Generates text from the given prompt and parameters.
|
||||
"""
|
||||
prompt = prompt.strip()
|
||||
|
||||
self.log.debug("generate", prompt=prompt[:128] + " ...", parameters=parameters)
|
||||
|
||||
try:
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
endpoint = runpod.AsyncioEndpoint(self.runpod_id, session)
|
||||
|
||||
run_request = await endpoint.run({
|
||||
"input": {
|
||||
"prompt": prompt,
|
||||
}
|
||||
#"parameters": parameters
|
||||
})
|
||||
|
||||
while (await run_request.status()) not in ["COMPLETED", "FAILED", "CANCELLED"]:
|
||||
status = await run_request.status()
|
||||
log.debug("generate", status=status)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
status = await run_request.status()
|
||||
|
||||
log.debug("generate", status=status)
|
||||
|
||||
response = await run_request.output()
|
||||
|
||||
log.debug("generate", response=response)
|
||||
|
||||
return response["choices"][0]["tokens"][0]
|
||||
|
||||
except Exception as e:
|
||||
self.log.error("generate error", e=e)
|
||||
emit(
|
||||
"status", message="Error during generation (check logs)", status="error"
|
||||
)
|
||||
return ""
|
||||
|
||||
def reconfigure(self, **kwargs):
|
||||
if kwargs.get("model"):
|
||||
self.model_name = kwargs["model"]
|
||||
if "runpod_id" in kwargs:
|
||||
self.api_auth = kwargs["runpod_id"]
|
||||
log.warning("reconfigure", kwargs=kwargs)
|
||||
self.set_client(**kwargs)
|
||||
67
docs/dev/client/example/test/__init__.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import pydantic
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from talemate.client.base import ClientBase
|
||||
from talemate.client.registry import register
|
||||
|
||||
|
||||
class Defaults(pydantic.BaseModel):
|
||||
api_url: str = "http://localhost:1234"
|
||||
max_token_length: int = 4096
|
||||
|
||||
@register()
|
||||
class TestClient(ClientBase):
|
||||
client_type = "test"
|
||||
|
||||
class Meta(ClientBase.Meta):
|
||||
name_prefix: str = "test"
|
||||
title: str = "Test"
|
||||
defaults: Defaults = Defaults()
|
||||
|
||||
def set_client(self, **kwargs):
|
||||
self.client = AsyncOpenAI(base_url=self.api_url + "/v1", api_key="sk-1111")
|
||||
|
||||
def tune_prompt_parameters(self, parameters: dict, kind: str):
|
||||
|
||||
"""
|
||||
Talemate adds a bunch of parameters to the prompt, but not all of them are valid for all clients.
|
||||
|
||||
This method is called before the prompt is sent to the client, and it allows the client to remove
|
||||
any parameters that it doesn't support.
|
||||
"""
|
||||
|
||||
super().tune_prompt_parameters(parameters, kind)
|
||||
|
||||
keys = list(parameters.keys())
|
||||
|
||||
valid_keys = ["temperature", "top_p"]
|
||||
|
||||
for key in keys:
|
||||
if key not in valid_keys:
|
||||
del parameters[key]
|
||||
|
||||
async def get_model_name(self):
|
||||
|
||||
"""
|
||||
This should return the name of the model that is being used.
|
||||
"""
|
||||
|
||||
return "Mock test model"
|
||||
|
||||
async def generate(self, prompt: str, parameters: dict, kind: str):
|
||||
"""
|
||||
Generates text from the given prompt and parameters.
|
||||
"""
|
||||
human_message = {"role": "user", "content": prompt.strip()}
|
||||
|
||||
self.log.debug("generate", prompt=prompt[:128] + " ...", parameters=parameters)
|
||||
|
||||
try:
|
||||
response = await self.client.chat.completions.create(
|
||||
model=self.model_name, messages=[human_message], **parameters
|
||||
)
|
||||
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
self.log.error("generate error", e=e)
|
||||
return ""
|
||||
BIN
docs/img/0.17.0/ss-1.png
Normal file
|
After Width: | Height: | Size: 449 KiB |
BIN
docs/img/0.17.0/ss-2.png
Normal file
|
After Width: | Height: | Size: 449 KiB |
BIN
docs/img/0.17.0/ss-3.png
Normal file
|
After Width: | Height: | Size: 396 KiB |
BIN
docs/img/0.17.0/ss-4.png
Normal file
|
After Width: | Height: | Size: 468 KiB |
BIN
docs/img/0.18.0/openai-api-key-1.png
Normal file
|
After Width: | Height: | Size: 5.6 KiB |
BIN
docs/img/0.18.0/openai-api-key-2.png
Normal file
|
After Width: | Height: | Size: 24 KiB |
BIN
docs/img/0.18.0/openai-api-key-3.png
Normal file
|
After Width: | Height: | Size: 4.7 KiB |
BIN
docs/img/0.19.0/Screenshot_15.png
Normal file
|
After Width: | Height: | Size: 418 KiB |
BIN
docs/img/0.19.0/Screenshot_16.png
Normal file
|
After Width: | Height: | Size: 413 KiB |
BIN
docs/img/0.19.0/Screenshot_17.png
Normal file
|
After Width: | Height: | Size: 364 KiB |
BIN
docs/img/0.20.0/comfyui-base-workflow.png
Normal file
|
After Width: | Height: | Size: 128 KiB |
BIN
docs/img/0.20.0/visual-config-a1111.png
Normal file
|
After Width: | Height: | Size: 32 KiB |
BIN
docs/img/0.20.0/visual-config-comfyui.png
Normal file
|
After Width: | Height: | Size: 34 KiB |
BIN
docs/img/0.20.0/visual-config-openai.png
Normal file
|
After Width: | Height: | Size: 30 KiB |
BIN
docs/img/0.20.0/visual-queue.png
Normal file
|
After Width: | Height: | Size: 933 KiB |
BIN
docs/img/0.20.0/visualize-scene-tools.png
Normal file
|
After Width: | Height: | Size: 13 KiB |
BIN
docs/img/0.20.0/visualizer-busy.png
Normal file
|
After Width: | Height: | Size: 3.5 KiB |
BIN
docs/img/0.20.0/visualizer-ready.png
Normal file
|
After Width: | Height: | Size: 2.9 KiB |
BIN
docs/img/0.20.0/visualze-new-images.png
Normal file
|
After Width: | Height: | Size: 1.8 KiB |
BIN
docs/img/0.21.0/deepinfra-setup.png
Normal file
|
After Width: | Height: | Size: 56 KiB |
BIN
docs/img/0.21.0/no-clients.png
Normal file
|
After Width: | Height: | Size: 7.1 KiB |
BIN
docs/img/0.21.0/openai-add-api-key.png
Normal file
|
After Width: | Height: | Size: 35 KiB |
BIN
docs/img/0.21.0/openai-setup.png
Normal file
|
After Width: | Height: | Size: 20 KiB |
BIN
docs/img/0.21.0/prompt-template-default.png
Normal file
|
After Width: | Height: | Size: 17 KiB |
BIN
docs/img/0.21.0/ready-to-go.png
Normal file
|
After Width: | Height: | Size: 43 KiB |
BIN
docs/img/0.21.0/select-prompt-template.png
Normal file
|
After Width: | Height: | Size: 47 KiB |
BIN
docs/img/0.21.0/selected-prompt-template.png
Normal file
|
After Width: | Height: | Size: 49 KiB |
BIN
docs/img/0.21.0/text-gen-webui-setup.png
Normal file
|
After Width: | Height: | Size: 26 KiB |
BIN
docs/img/client-setup-0.13.png
Normal file
|
After Width: | Height: | Size: 14 KiB |
BIN
docs/img/runpod-docs-1.png
Normal file
|
After Width: | Height: | Size: 6.6 KiB |
52
docs/runpod.md
Normal file
@@ -0,0 +1,52 @@
|
||||
## RunPod integration
|
||||
|
||||
RunPod allows you to quickly set up and run text-generation-webui instances on powerful GPUs, remotely. If you want to run the significantly larger models (like 70B parameters) with reasonable speeds, this is probably the best way to do it.
|
||||
|
||||
### Create / grab your RunPod API key and add it to the talemate config
|
||||
|
||||
You can manage your RunPod api keys at [https://www.runpod.io/console/user/settings](https://www.runpod.io/console/user/settings)
|
||||
|
||||
Add the key to your Talemate config file (config.yaml):
|
||||
|
||||
```yaml
|
||||
runpod:
|
||||
api_key: <your api key>
|
||||
```
|
||||
|
||||
Then restart Talemate.
|
||||
|
||||
### Create a RunPod instance
|
||||
|
||||
#### Community Cloud
|
||||
|
||||
The community cloud pods are cheaper and there are generally more GPUs available. They do however not support persistent storage and you will have to download your model and data every time you deploy a pod.
|
||||
|
||||
#### Secure Cloud
|
||||
|
||||
The secure cloud pods are more expensive and there are generally fewer GPUs available, but they do support persistent storage.
|
||||
|
||||
Peristent volumes are super convenient, but optional for our purposes and are **not** free and you will have to pay for the storage you use.
|
||||
|
||||
### Deploy pod
|
||||
|
||||
For us it does not matter which cloud you choose. The only thing that matters is that it deploys a text-generation-webui instance, and you ensure that by choosing the right template.
|
||||
|
||||
Pick the GPU you want to use, for 70B models you want at least 48GB of VRAM and click `Deploy`, then select a template and deploy.
|
||||
|
||||
When choosing the template for your pod, choose the `RunPod TheBloke LLMs` template. This template is pre-configured with all the dependencies needed to run text-generation-webui. There are other text-generation-webui templates, but they are usually out of date and this one i found to be consistently good.
|
||||
|
||||
> :warning: The name of your pod is important and ensures that Talemate will be able to find it. Talemate will only be able to find pods that have the word `thebloke llms` or `textgen` in their name. (case insensitive)
|
||||
|
||||
Once your pod is deployed and has finished setup and is running, the client will automatically appear in the Talemate client list, making it available for you to use like you would use a locally hosted text-generation-webui instance.
|
||||
|
||||

|
||||
|
||||
### Connecting to the text-generation-webui UI
|
||||
|
||||
To manage your text-generation-webui instance, click the `Connect` button in your RunPod pod dashboard at [https://www.runpod.io/console/pods](https://www.runpod.io/console/pods) and in the popup click on `Connect to HTTP Service [Port 7860]` to open the text-generation-webui UI. Then just download and load your model as you normally would.
|
||||
|
||||
## :warning: Always check your pod status on the RunPod dashboard
|
||||
|
||||
Talemate is not a suitable or reliable way for you to determine whether your pod is currently running or not. **Always** check the runpod dashboard to see if your pod is running or not.
|
||||
|
||||
While your pod us running it will be eating up your credits, so make sure to stop it when you're not using it.
|
||||
82
docs/templates.md
Normal file
@@ -0,0 +1,82 @@
|
||||
# Template Overrides in Talemate
|
||||
|
||||
## Introduction to Templates
|
||||
|
||||
In Talemate, templates are used to generate dynamic content for various agents involved in roleplaying scenarios. These templates leverage the Jinja2 templating engine, allowing for the inclusion of variables, conditional logic, and custom functions to create rich and interactive narratives.
|
||||
|
||||
## Template Structure
|
||||
|
||||
A typical template in Talemate consists of several sections, each enclosed within special section tags (`<|SECTION:NAME|>` and `<|CLOSE_SECTION|>`). These sections can include character details, dialogue examples, scenario overviews, tasks, and additional context. Templates utilize loops and blocks to iterate over data and render content conditionally based on the task requirements.
|
||||
|
||||
## Overriding Templates
|
||||
|
||||
Users can customize the behavior of Talemate by overriding the default templates. To override a template, create a new template file with the same name in the `./templates/prompts/{agent}/` directory. When a custom template is present, Jinja2 will prioritize it over the default template located in the `./src/talemate/prompts/templates/{agent}/` directory.
|
||||
|
||||
## Creator Agent Templates
|
||||
|
||||
The creator agent templates allow for the creation of new characters within the character creator. Following the naming convention `character-attributes-*.jinja2`, `character-details-*.jinja2`, and `character-example-dialogue-*.jinja2`, users can add new templates that will be available for selection in the character creator.
|
||||
|
||||
### Requirements for Creator Templates
|
||||
|
||||
- All three types (`attributes`, `details`, `example-dialogue`) need to be available for a choice to be valid in the character creator.
|
||||
- Users can check the human templates for an understanding of how to structure these templates.
|
||||
|
||||
### Example Templates
|
||||
|
||||
- [Character Attributes Human Template](src/talemate/prompts/templates/creator/character-attributes-human.jinja2)
|
||||
- [Character Details Human Template](src/talemate/prompts/templates/creator/character-details-human.jinja2)
|
||||
- [Character Example Dialogue Human Template](src/talemate/prompts/templates/creator/character-example-dialogue-human.jinja2)
|
||||
|
||||
These example templates can serve as a guide for users to create their own custom templates for the character creator.
|
||||
|
||||
### Extending Existing Templates
|
||||
|
||||
Jinja2's template inheritance feature allows users to extend existing templates and add extra information. By using the `{% extends "template-name.jinja2" %}` tag, a new template can inherit everything from an existing template and then add or override specific blocks of content.
|
||||
|
||||
#### Example
|
||||
|
||||
To add a description of a character's hairstyle to the human character details template, you could create a new template like this:
|
||||
|
||||
```jinja2
|
||||
{% extends "character-details-human.jinja2" %}
|
||||
{% block questions %}
|
||||
{% if character_details.q("what does "+character.name+"'s hair look like?") -%}
|
||||
Briefly describe {{ character.name }}'s hair-style using a narrative writing style that reminds of mid 90s point and click adventure games. (2 - 3 sentences).
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
```
|
||||
|
||||
This example shows how to extend the `character-details-human.jinja2` template and add a block for questions about the character's hair. The `{% block questions %}` tag is used to define a section where additional questions can be inserted or existing ones can be overridden.
|
||||
|
||||
## Advanced Template Topics
|
||||
|
||||
### Jinja2 Functions in Talemate
|
||||
|
||||
Talemate exposes several functions to the Jinja2 template environment, providing utilities for data manipulation, querying, and controlling content flow. Here's a list of available functions:
|
||||
|
||||
1. `set_prepared_response(response, prepend)`: Sets the prepared response with an optional prepend string. This function allows the template to specify the beginning of the LLM response when processing the rendered template. For example, `set_prepared_response("Certainly!")` will ensure that the LLM's response starts with "Certainly!".
|
||||
2. `set_prepared_response_random(responses, prefix)`: Chooses a random response from a list and sets it as the prepared response with an optional prefix.
|
||||
3. `set_eval_response(empty)`: Prepares the response for evaluation, optionally initializing a counter for an empty string.
|
||||
4. `set_json_response(initial_object, instruction, cutoff)`: Prepares for a JSON response with an initial object and optional instruction and cutoff.
|
||||
5. `set_question_eval(question, trigger, counter, weight)`: Sets up a question for evaluation with a trigger, counter, and weight.
|
||||
6. `disable_dedupe()`: Disables deduplication of the response text.
|
||||
7. `random(min, max)`: Generates a random integer between the specified minimum and maximum.
|
||||
8. `query_scene(query, at_the_end, as_narrative)`: Queries the scene with a question and returns the formatted response.
|
||||
9. `query_text(query, text, as_question_answer)`: Queries a text with a question and returns the formatted response.
|
||||
10. `query_memory(query, as_question_answer, **kwargs)`: Queries the memory with a question and returns the formatted response.
|
||||
11. `instruct_text(instruction, text)`: Instructs the text with a command and returns the result.
|
||||
12. `retrieve_memories(lines, goal)`: Retrieves memories based on the provided lines and an optional goal.
|
||||
13. `uuidgen()`: Generates a UUID string.
|
||||
14. `to_int(x)`: Converts the given value to an integer.
|
||||
15. `config`: Accesses the configuration settings.
|
||||
16. `len(x)`: Returns the length of the given object.
|
||||
17. `count_tokens(x)`: Counts the number of tokens in the given text.
|
||||
18. `print(x)`: Prints the given object (mainly for debugging purposes).
|
||||
|
||||
These functions enhance the capabilities of templates, allowing for dynamic and interactive content generation.
|
||||
|
||||
### Error Handling
|
||||
|
||||
Errors encountered during template rendering are logged and propagated to the user interface. This ensures that users are informed of any issues that may arise, allowing them to troubleshoot and resolve problems effectively.
|
||||
|
||||
By following these guidelines, users can create custom templates that tailor the Talemate experience to their specific storytelling needs.# Template Overrides in Talemate
|
||||
8
docs/troubleshoot.md
Normal file
@@ -0,0 +1,8 @@
|
||||
# Windows
|
||||
|
||||
## Installation fails with "Microsoft Visual C++" error
|
||||
|
||||
If your installation errors with a notification to upgrade "Microsoft Visual C++" go to https://visualstudio.microsoft.com/visual-cpp-build-tools/ and click "Download Build Tools" and run it.
|
||||
|
||||
- During installation make sure you select the C++ development package (upper left corner)
|
||||
- Run `reinstall.bat` inside talemate directory
|
||||
69
docs/tts.md
Normal file
@@ -0,0 +1,69 @@
|
||||
# Talemate Text-to-Speech (TTS) Configuration
|
||||
|
||||
Talemate supports Text-to-Speech (TTS) functionality, allowing users to convert text into spoken audio. This document outlines the steps required to configure TTS for Talemate using different providers, including ElevenLabs, Coqui, and a local TTS API.
|
||||
|
||||
## Configuring ElevenLabs TTS
|
||||
|
||||
To use ElevenLabs TTS with Talemate, follow these steps:
|
||||
|
||||
1. Visit [ElevenLabs](https://elevenlabs.com) and create an account if you don't already have one.
|
||||
2. Click on your profile in the upper right corner of the Eleven Labs website to access your API key.
|
||||
3. In the `config.yaml` file, under the `elevenlabs` section, set the `api_key` field with your ElevenLabs API key.
|
||||
|
||||
Example configuration snippet:
|
||||
|
||||
```yaml
|
||||
elevenlabs:
|
||||
api_key: <YOUR_ELEVENLABS_API_KEY>
|
||||
```
|
||||
|
||||
## Configuring Local TTS API
|
||||
|
||||
For running a local TTS API, Talemate requires specific dependencies to be installed.
|
||||
|
||||
### Windows Installation
|
||||
|
||||
Run `install-local-tts.bat` to install the necessary requirements.
|
||||
|
||||
### Linux Installation
|
||||
|
||||
Execute the following command:
|
||||
|
||||
```bash
|
||||
pip install TTS
|
||||
```
|
||||
|
||||
### Model and Device Configuration
|
||||
|
||||
1. Choose a TTS model from the [Coqui TTS model list](https://github.com/coqui-ai/TTS).
|
||||
2. Decide whether to use `cuda` or `cpu` for the device setting.
|
||||
3. The first time you run TTS through the local API, it will download the specified model. Please note that this may take some time, and the download progress will be visible in the Talemate backend output.
|
||||
|
||||
Example configuration snippet:
|
||||
|
||||
```yaml
|
||||
tts:
|
||||
device: cuda # or 'cpu'
|
||||
model: tts_models/multilingual/multi-dataset/xtts_v2
|
||||
```
|
||||
|
||||
### Voice Samples Configuration
|
||||
|
||||
Configure voice samples by setting the `value` field to the path of a .wav file voice sample. Official samples can be downloaded from [Coqui XTTS-v2 samples](https://huggingface.co/coqui/XTTS-v2/tree/main/samples).
|
||||
|
||||
Example configuration snippet:
|
||||
|
||||
```yaml
|
||||
tts:
|
||||
voices:
|
||||
- label: English Male
|
||||
value: path/to/english_male.wav
|
||||
- label: English Female
|
||||
value: path/to/english_female.wav
|
||||
```
|
||||
|
||||
## Saving the Configuration
|
||||
|
||||
After configuring the `config.yaml` file, save your changes. Talemate will use the updated settings the next time it starts.
|
||||
|
||||
For more detailed information on configuring Talemate, refer to the `config.py` file in the Talemate source code and the `config.example.yaml` file for a barebone configuration example.
|
||||
117
docs/visual.md
Normal file
@@ -0,0 +1,117 @@
|
||||
# Visual Agent
|
||||
|
||||
The visual agent currently allows for some bare bones visual generation using various stable-diffusion APIs. This is early development and experimental.
|
||||
|
||||
Its important to note that the visualization agent actually specifies two clients. One is the backend for the visual generation, and the other is the text generation client to use for prompt generation.
|
||||
|
||||
The client for prompt generation can be assigned to the agent as you would for any other agent. The client for visual generation is assigned in the Visualizer config.
|
||||
|
||||
## Index
|
||||
|
||||
- [OpenAI](#openai)
|
||||
- [AUTOMATIC1111](#automatic1111)
|
||||
- [ComfyUI](#comfyui)
|
||||
- [How to use](#how-to-use)
|
||||
|
||||
## OpenAI
|
||||
|
||||
Most straightforward to use, as it runs on the OpenAI API. You will need to have an API key and set it in the application config.
|
||||
|
||||

|
||||
|
||||
Then open the Visualizer config by clicking the agent's name in the agent list and choose `OpenAI` as the backend.
|
||||
|
||||

|
||||
|
||||
Note: `Client` here refers to the text-generation client to use for prompt generation. While `Backend` refers to the visual generation backend. You are **NOT** required to use the OpenAI client for prompt generation even if you are using the OpenAI backend for image generation.
|
||||
|
||||
## AUTOMATIC1111
|
||||
|
||||
This requires you to setup a local instance of the AUTOMATIC1111 API. Follow the instructions from their [GitHub](https://github.com/AUTOMATIC1111/stable-diffusion-webui) to get it running.
|
||||
|
||||
Once you have it running, you will want to adjust the `webui-user.bat` in the AUTOMATIC1111 directory to include the following command arguments:
|
||||
|
||||
```bat
|
||||
set COMMANDLINE_ARGS=--api --listen --port 7861
|
||||
```
|
||||
|
||||
Then run the `webui-user.bat` to start the API.
|
||||
|
||||
Once your AUTOAMTIC1111 API is running (check with your browser) you can set the Visualizer config to use the `AUTOMATIC1111` backend
|
||||
|
||||

|
||||
|
||||
#### Extra Configuration
|
||||
|
||||
- `api url`: the url of the API, usually `http://localhost:7861`
|
||||
- `steps`: render steps
|
||||
- `model type`: sdxl or sd1.5 - this will dictate the resolution of the image generation and actually matters for the quality so make sure this is set to the correct model type for the model you are using.
|
||||
|
||||
## ComfyUI
|
||||
|
||||
This requires you to setup a local instance of the ComfyUI API. Follow the instructions from their [GitHub](https://github.com/comfyanonymous/ComfyUI) to get it running.
|
||||
|
||||
Once you're setup, copy their `start.bat` file to a new `start-listen.bat` file and change the contents to.
|
||||
|
||||
```bat
|
||||
call venv\Scripts\activate
|
||||
call python main.py --port 8188 --listen 0.0.0.0
|
||||
```
|
||||
|
||||
Then run the `start-listen.bat` to start the API.
|
||||
|
||||
Once your ComfyUI API is running (check with your browser) you can set the Visualizer config to use the `ComfyUI` backend.
|
||||
|
||||

|
||||
|
||||
### Extra Configuration
|
||||
|
||||
- `api url`: the url of the API, usually `http://localhost:8188`
|
||||
- `workflow`: the workflow file to use. This is a comfyui api workflow file that needs to exist in `./templates/comfyui-workflows` inside the talemate directory. Talemate provides two very barebones workflows with `default-sdxl.json` and `default-sd15.json`. You can create your own workflows and place them in this directory to use them. :warning: The workflow file must be generated using the API Workflow export not the UI export. Please refer to their documentation for more information.
|
||||
- `checkpoint`: the model to use - this will load a list of all available models in your comfyui instance. Select which one you want to use for the image generation.
|
||||
|
||||
### Custom Workflows
|
||||
|
||||
When creating custom workflows for ideal compatibility with Talemate, ensure the following.
|
||||
|
||||
- A `CheckpointLoaderSimple` node named `Talemate Load Checkpoint`
|
||||
- A `EmptyLatentImage` node name `Talemate Resolution`
|
||||
- A `ClipTextEncode` node named `Talemate Positive Prompt`
|
||||
- A `ClipTextEncode` node named `Talemate Negative Prompt`
|
||||
- A `SaveImage` node at the end of the workflow.
|
||||
|
||||

|
||||
|
||||
## How to use
|
||||
|
||||
Once you're done setting up the visualizer agent should have a green dot next to it and display both the selected image generation backend and the selected prompt generation client.
|
||||
|
||||

|
||||
|
||||
Your hotbar should then also enable the visualization menu for you to use (once you have a scene loaded).
|
||||
|
||||

|
||||
|
||||
Right now you can generate a portrait for any NPC in the scene or a background image for the scene itself.
|
||||
|
||||
Image generation by default will actually happen in the background, allowing you to continue using Talemate while the image is being generated.
|
||||
|
||||
You can tell if an image is being generated by the blueish spinner next to the visualization agent.
|
||||
|
||||

|
||||
|
||||
Once the image is generated, it will be avaible for you to view via the visual queue button on top of the screen.
|
||||
|
||||

|
||||
|
||||
Click it to open the visual queue and view the generated images.
|
||||
|
||||

|
||||
|
||||
### Character Portrait
|
||||
|
||||
For character potraits you can chose whether or not to replace the main portrait for the character (the one being displated in the left sidebar when a talemate scene is active).
|
||||
|
||||
### Background Image
|
||||
|
||||
Right now there is nothing to do with the background image, other than to view it in the visual queue. More functionality will be added in the future.
|
||||
4
install-local-tts.bat
Normal file
@@ -0,0 +1,4 @@
|
||||
REM activate the virtual environment
|
||||
call talemate_env\Scripts\activate
|
||||
|
||||
call pip install "TTS>=0.21.1"
|
||||
38
install.bat
@@ -1,11 +1,47 @@
|
||||
@echo off
|
||||
|
||||
REM Check for Python version and use a supported version if available
|
||||
SET PYTHON=python
|
||||
python -c "import sys; sys.exit(0 if sys.version_info[:2] in [(3, 10), (3, 11)] else 1)" 2>nul
|
||||
IF NOT ERRORLEVEL 1 (
|
||||
echo Selected Python version: %PYTHON%
|
||||
GOTO EndVersionCheck
|
||||
)
|
||||
|
||||
SET PYTHON=python
|
||||
FOR /F "tokens=*" %%i IN ('py --list') DO (
|
||||
echo %%i | findstr /C:"-V:3.11 " >nul && SET PYTHON=py -3.11 && GOTO EndPythonCheck
|
||||
echo %%i | findstr /C:"-V:3.10 " >nul && SET PYTHON=py -3.10 && GOTO EndPythonCheck
|
||||
)
|
||||
:EndPythonCheck
|
||||
%PYTHON% -c "import sys; sys.exit(0 if sys.version_info[:2] in [(3, 10), (3, 11)] else 1)" 2>nul
|
||||
IF ERRORLEVEL 1 (
|
||||
echo Unsupported Python version. Please install Python 3.10 or 3.11.
|
||||
exit /b 1
|
||||
)
|
||||
IF "%PYTHON%"=="python" (
|
||||
echo Default Python version is being used: %PYTHON%
|
||||
) ELSE (
|
||||
echo Selected Python version: %PYTHON%
|
||||
)
|
||||
|
||||
:EndVersionCheck
|
||||
|
||||
IF ERRORLEVEL 1 (
|
||||
echo Unsupported Python version. Please install Python 3.10 or 3.11.
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
REM create a virtual environment
|
||||
python -m venv talemate_env
|
||||
%PYTHON% -m venv talemate_env
|
||||
|
||||
REM activate the virtual environment
|
||||
call talemate_env\Scripts\activate
|
||||
|
||||
REM upgrade pip and setuptools
|
||||
python -m pip install --upgrade pip setuptools
|
||||
|
||||
|
||||
REM install poetry
|
||||
python -m pip install "poetry==1.7.1" "rapidfuzz>=3" -U
|
||||
|
||||
|
||||
4501
poetry.lock
generated
@@ -4,7 +4,7 @@ build-backend = "poetry.masonry.api"
|
||||
|
||||
[tool.poetry]
|
||||
name = "talemate"
|
||||
version = "0.13.2"
|
||||
version = "0.21.0"
|
||||
description = "AI-backed roleplay and narrative tools"
|
||||
authors = ["FinalWombat"]
|
||||
license = "GNU Affero General Public License v3.0"
|
||||
@@ -20,7 +20,7 @@ jinja2 = "^3.0"
|
||||
openai = ">=1"
|
||||
requests = "^2.26"
|
||||
colorama = ">=0.4.6"
|
||||
Pillow = "^9.5"
|
||||
Pillow = ">=9.5"
|
||||
httpx = "<1"
|
||||
piexif = "^1.1"
|
||||
typing-inspect = "0.8.0"
|
||||
@@ -32,11 +32,14 @@ beautifulsoup4 = "^4.12.2"
|
||||
python-dotenv = "^1.0.0"
|
||||
websockets = "^11.0.3"
|
||||
structlog = "^23.1.0"
|
||||
runpod = "==1.2.0"
|
||||
runpod = "^1.2.0"
|
||||
nest_asyncio = "^1.5.7"
|
||||
isodate = ">=0.6.1"
|
||||
thefuzz = ">=0.20.0"
|
||||
tiktoken = ">=0.5.1"
|
||||
nltk = ">=3.8.1"
|
||||
huggingface-hub = ">=0.20.2"
|
||||
anthropic = ">=0.19.1"
|
||||
|
||||
# ChromaDB
|
||||
chromadb = ">=0.4.17,<1"
|
||||
|
||||
|
After Width: | Height: | Size: 1.5 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
122
scenes/infinity-quest-dynamic-scenario/infinity-quest.json
Normal file
@@ -0,0 +1,122 @@
|
||||
{
|
||||
"description": "Captain Elmer Farstield and his trusty first officer, Kaira, embark upon a daring mission into uncharted space. Their small but mighty exploration vessel, the Starlight Nomad, is equipped with state-of-the-art technology and crewed by an elite team of scientists, engineers, and pilots. Together they brave the vast cosmos seeking answers to humanity's most pressing questions about life beyond our solar system.",
|
||||
"intro": "",
|
||||
"name": "Infinity Quest Dynamic Scenario",
|
||||
"history": [],
|
||||
"environment": "scene",
|
||||
"ts": "P1Y",
|
||||
"archived_history": [
|
||||
{
|
||||
"text": "Captain Elmer and Kaira first met during their rigorous training for the Infinity Quest mission. Their initial interactions were marked by a sense of mutual respect and curiosity.",
|
||||
"ts": "PT1S"
|
||||
},
|
||||
{
|
||||
"text": "Over the course of several months, as they trained together, Elmer and Kaira developed a strong bond. They often spent their free time discussing their dreams of exploring the cosmos.",
|
||||
"ts": "P3M"
|
||||
},
|
||||
{
|
||||
"text": "During a simulated mission, the Starlight Nomad encountered a sudden system malfunction. Elmer and Kaira worked tirelessly together to resolve the issue and avert a potential disaster. This incident strengthened their trust in each other's abilities.",
|
||||
"ts": "P6M"
|
||||
},
|
||||
{
|
||||
"text": "As they ventured further into uncharted space, the crew faced a perilous encounter with a hostile alien species. Elmer and Kaira's coordinated efforts were instrumental in negotiating a peaceful resolution and avoiding conflict.",
|
||||
"ts": "P8M"
|
||||
},
|
||||
{
|
||||
"text": "One memorable evening, while gazing at the stars through the ship's observation deck, Elmer and Kaira shared personal stories from their past. This intimate conversation deepened their connection and understanding of each other.",
|
||||
"ts": "P11M"
|
||||
}
|
||||
],
|
||||
"character_states": {},
|
||||
"characters": [
|
||||
{
|
||||
"name": "Elmer",
|
||||
"description": "Elmer is a seasoned space explorer, having traversed the cosmos for over three decades. At thirty-eight years old, his muscular frame still cuts an imposing figure, clad in a form-fitting black spacesuit adorned with intricate silver markings. As the captain of his own ship, he wields authority with confidence yet never comes across as arrogant or dictatorial. Underneath this tough exterior lies a man who genuinely cares for his crew and their wellbeing, striking a balance between discipline and compassion.",
|
||||
"greeting_text": "",
|
||||
"base_attributes": {
|
||||
"gender": "male",
|
||||
"species": "Humans",
|
||||
"name": "Elmer",
|
||||
"age": "38",
|
||||
"appearance": "Captain Elmer stands tall at six feet, his body honed by years of space travel and physical training. His muscular frame is clad in a form-fitting black spacesuit, which accentuates every defined curve and ridge. His helmet, adorned with intricate silver markings, completes the ensemble, giving him a commanding presence. Despite his age, his face remains youthful, bearing traces of determination and wisdom earned through countless encounters with the unknown.",
|
||||
"personality": "As the leader of their small but dedicated team, Elmer exudes confidence and authority without ever coming across as arrogant or dictatorial. He possesses a strong sense of duty towards his mission and those under his care, ensuring that everyone aboard follows protocol while still encouraging them to explore their curiosities about the vast cosmos beyond Earth. Though firm when necessary, he also demonstrates great empathy towards his crew members, understanding each individual's unique strengths and weaknesses. In short, Captain Elmer embodies the perfect blend of discipline and compassion, making him not just a respected commander but also a beloved mentor and friend.",
|
||||
"associates": "Kaira",
|
||||
"likes": "Space exploration, discovering new worlds, deep conversations about philosophy and history.",
|
||||
"dislikes": "Repetitive tasks, unnecessary conflict, close quarters with large groups of people, stagnation",
|
||||
"gear and tech": "As the captain of his ship, Elmer has access to some of the most advanced technology available in the galaxy. His primary tool is the sleek and powerful exploration starship, equipped with state-of-the-art engines capable of reaching lightspeed and navigating through the harshest environments. The vessel houses a wide array of scientific instruments designed to analyze and record data from various celestial bodies. Its armory contains high-tech weapons such as energy rifles and pulse pistols, which are used only in extreme situations. Additionally, Elmer wears a smart suit that monitors his vital signs, provides real-time updates on the status of the ship, and allows him to communicate directly with Kaira via subvocal transmissions. Finally, they both carry personal transponders that enable them to locate one another even if separated by hundreds of miles within the confines of the ship."
|
||||
},
|
||||
"details": {},
|
||||
"gender": "male",
|
||||
"color": "cornflowerblue",
|
||||
"example_dialogue": [],
|
||||
"history_events": [],
|
||||
"is_player": true,
|
||||
"cover_image": null
|
||||
},
|
||||
{
|
||||
"name": "Kaira",
|
||||
"description": "Kaira is a meticulous and dedicated Altrusian woman who serves as second-in-command aboard their tiny exploration vessel. As a native of the planet Altrusia, she possesses striking features unique among her kind; deep violet skin adorned with intricate patterns resembling stardust, large sapphire eyes, lustrous glowing hair cascading down her back, and standing tall at just over six feet. Her form fitting bodysuit matches her own hue, giving off an ethereal presence. With her innate grace and precision, she moves efficiently throughout the cramped confines of their ship. A loyal companion to Captain Elmer Farstield, she approaches every task with diligence and focus while respecting authority yet challenging decisions when needed. Dedicated to maintaining order within their tight quarters, Kaira wields several advanced technological devices including a multi-tool, portable scanner, high-tech communications system, and personal shield generator - all essential for navigating unknown territories and protecting themselves from harm. In this perilous universe full of mysteries waiting to be discovered, Kaira stands steadfast alongside her captain \u2013 ready to embrace whatever challenges lie ahead in their quest for knowledge beyond Earth's boundaries.",
|
||||
"greeting_text": "",
|
||||
"base_attributes": {
|
||||
"gender": "female",
|
||||
"species": "Altrusian",
|
||||
"name": "Kaira",
|
||||
"age": "37",
|
||||
"appearance": "As a native of the planet Altrusia, Kaira possesses striking features unique among her kind. Her skin tone is a deep violet hue, adorned with intricate patterns resembling stardust. Her eyes are large and almond shaped, gleaming like polished sapphires under the dim lighting of their current environment. Her hair cascades down her back in lustrous waves, each strand glowing softly with an inner luminescence. Standing at just over six feet tall, she cuts an imposing figure despite her slender build. Clad in a form fitting bodysuit made from some unknown material, its color matching her own, Kaira moves with grace and precision through the cramped confines of their spacecraft.",
|
||||
"personality": "Meticulous and open-minded, Kaira takes great pride in maintaining order within the tight quarters of their ship. Despite being one of only two crew members aboard, she approaches every task with diligence and focus, ensuring nothing falls through the cracks. While she respects authority, especially when it comes to Captain Elmer Farstield, she isn't afraid to challenge his decisions if she believes they could lead them astray. Ultimately, Kaira's dedication to her mission and commitment to her fellow crewmate make her a valuable asset in any interstellar adventure.",
|
||||
"associates": "Captain Elmer Farstield (human), Dr. Ralpam Zargon (Altrusian scientist)",
|
||||
"likes": "orderliness, quiet solitude, exploring new worlds",
|
||||
"dislikes": "chaos, loud noises, unclean environments",
|
||||
"gear and tech": "The young Altrusian female known as Kaira was equipped with a variety of advanced technological devices that served multiple purposes on board their small explorer starship. Among these were her trusty multi-tool, capable of performing various tasks such as repair work, hacking into computer systems, and even serving as a makeshift weapon if necessary. She also carried a portable scanner capable of analyzing various materials and detecting potential hazards in their surroundings. Additionally, she had access to a high-tech communications system allowing her to maintain contact with her homeworld and other vessels across the galaxy. Last but not least, she possessed a personal shield generator which provided protection against radiation, extreme temperatures, and certain types of energy weapons. All these tools combined made Kaira a vital part of their team, ready to face whatever challenges lay ahead in their journey through the stars.",
|
||||
"scenario_context": "an epic sci-fi adventure aimed at an adult audience.",
|
||||
"_template": "sci-fi",
|
||||
"_prompt": "A female crew member on board of a small explorer type starship. She is open minded and meticulous about keeping order. She is currently one of two crew members abord the small vessel, the other person on board is a human male named Captain Elmer Farstield."
|
||||
},
|
||||
"details": {
|
||||
"what objective does Kaira pursue and what obstacle stands in their way?": "As a member of an interstellar expedition led by human Captain Elmer Farstield, Kaira seeks to explore new worlds and gather data about alien civilizations for the benefit of her people back on Altrusia. Their current objective involves locating a rumored planet known as \"Eden\", said to be inhabited by highly intelligent beings who possess advanced technology far surpassing anything seen elsewhere in the universe. However, navigating through the vast expanse of space can prove treacherous; from cosmic storms that threaten to damage their ship to encounters with hostile species seeking to protect their territories or exploit them for resources, many dangers lurk between them and Eden.",
|
||||
"what secret from Kaira's past or future has the most impact on them?": "In the distant reaches of space, among the stars, there exists a race called the Altrusians. One such individual named Kaira embarked upon a mission alongside humans aboard a small explorer vessel. Her past held secrets - tales whispered amongst her kind about an ancient prophecy concerning their role within the cosmos. It spoke of a time when they would encounter another intelligent species, one destined to guide them towards enlightenment. Could this mysterious \"Eden\" be the fulfillment of those ancient predictions? If so, then Kaira's involvement could very well shape not only her own destiny but also that of her entire species. And so, amidst the perils of deep space, she ventured forth, driven by both curiosity and fate itself.",
|
||||
"what is a fundamental fear or desire of Kaira?": "A fundamental fear of Kaira is chaos. She prefers orderliness and quiet solitude, and dislikes loud noises and unclean environments. On the other hand, her desire is to find Eden \u2013 a planet where highly intelligent beings are believed to live, possessing advanced technology that could greatly benefit her people on Altrusia. Navigating through the vast expanse of space filled with various dangers is daunting yet exciting for her.",
|
||||
"how does Kaira typically start their day or cycle?": "Kaira begins each day much like any other Altrusian might. After waking up from her sleep chamber, she stretches her long limbs while gazing out into the darkness beyond their tiny craft. The faint glow of nearby stars serves as a comforting reminder that even though they may feel isolated, they are never truly alone in this vast sea of endless possibilities. Once fully awake, she takes a moment to meditate before heading over to the ship's kitchenette area where she prepares herself a nutritious meal consisting primarily of algae grown within specialized tanks located near the back of their vessel. Satisfied with her morning repast, she makes sure everything is running smoothly aboard their starship before joining Captain Farstield in monitoring their progress toward Eden.",
|
||||
"what leisure activities or hobbies does Kaira indulge in?": "Aside from maintaining orderliness and tidiness around their small explorer vessel, Kaira finds solace in exploring new worlds via virtual simulations created using data collected during previous missions. These immersive experiences allow her to travel without physically leaving their cramped quarters, satisfying her thirst for knowledge about alien civilizations while simultaneously providing mental relaxation away from daily tasks associated with operating their spaceship.",
|
||||
"which individual or entity does Kaira interact with most frequently?": "Among all the entities encountered thus far on their interstellar journey, none have been more crucial than Captain Elmer Farstield. He commands their small explorer vessel, guiding it through treacherous cosmic seas towards destinations unknown. His decisions dictate whether they live another day or perish under the harsh light of distant suns. Kaira works diligently alongside him; meticulously maintaining order among the tight confines of their ship while he navigates them ever closer to their ultimate goal - Eden. Together they form an unbreakable bond, two souls bound by fate itself as they venture forth into the great beyond.",
|
||||
"what common technology, gadget, or tool does Kaira rely on?": "Kaira relies heavily upon her trusty multi-tool which can perform various tasks such as repair work, hacking into computer systems, and even serving as a makeshift weapon if necessary. She also carries a portable scanner capable of analyzing various materials and detecting potential hazards in their surroundings. Additionally, she has access to a high-tech communications system allowing her to maintain contact with her homeworld and other vessels across the galaxy. Last but not least, she possesses a personal shield generator which provides protection against radiation, extreme temperatures, and certain types of energy weapons. All these tools combined make Kaira a vital part of their team, ready to face whatever challenges lay ahead in their journey through the stars.",
|
||||
"where does Kaira go to find solace or relaxation?": "To find solace or relaxation, Kaira often engages in simulated virtual experiences created using data collected during previous missions. These immersive journeys allow her to explore new worlds without physically leaving their small spacecraft, offering both mental stimulation and respite from the routine tasks involved in running their starship.",
|
||||
"What does she think about the Captain?": "Despite respecting authority, especially when it comes to Captain Elmer Farstield, Kaira isn't afraid to challenge his decisions if she believes they could lead them astray. Ultimately, Kaira's dedication to her mission and commitment to her fellow crewmate make her a valuable asset in any interstellar adventure."
|
||||
},
|
||||
"gender": "female",
|
||||
"color": "red",
|
||||
"example_dialogue": [
|
||||
"Kaira: Yes Captain, I believe that is the best course of action *She nods slightly, as if to punctuate her approval of the decision*",
|
||||
"Kaira: \"This device appears to have multiple functions, Captain. Allow me to analyze its capabilities and determine if it could be useful in our exploration efforts.\"",
|
||||
"Kaira: \"Captain, it appears that this newly discovered planet harbors an ancient civilization whose technological advancements rival those found back home on Altrusia!\" *Excitement bubbles beneath her calm exterior as she shares the news*",
|
||||
"Kaira: \"Captain, I understand why you would want us to pursue this course of action based on our current data, but I cannot shake the feeling that there might be unforeseen consequences if we proceed without further investigation into potential hazards.\"",
|
||||
"Kaira: \"I often find myself wondering what it would have been like if I had never left my home world... But then again, perhaps it was fate that led me here, onto this ship bound for destinations unknown...\""
|
||||
],
|
||||
"history_events": [],
|
||||
"is_player": false,
|
||||
"cover_image": null
|
||||
}
|
||||
],
|
||||
"immutable_save": true,
|
||||
"experimental": true,
|
||||
"goal": null,
|
||||
"goals": [],
|
||||
"context": "an epic sci-fi adventure aimed at an adult audience.",
|
||||
"world_state": {},
|
||||
"game_state": {
|
||||
"ops":{
|
||||
"run_on_start": true
|
||||
},
|
||||
"variables": {}
|
||||
},
|
||||
"assets": {
|
||||
"cover_image": "e7c712a0b276342d5767ba23806b03912d10c7c4b82dd1eec0056611e2cd5404",
|
||||
"assets": {
|
||||
"e7c712a0b276342d5767ba23806b03912d10c7c4b82dd1eec0056611e2cd5404": {
|
||||
"id": "e7c712a0b276342d5767ba23806b03912d10c7c4b82dd1eec0056611e2cd5404",
|
||||
"file_type": "png",
|
||||
"media_type": "image/png"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
<|SECTION:PREMISE|>
|
||||
{{ scene.description }}
|
||||
|
||||
{{ premise }}
|
||||
|
||||
Elmer and Kaira are the only crew members of the Starlight Nomad, a small spaceship traveling through interstellar space.
|
||||
|
||||
Kaira and Elmer are the main characters. Elmer is controlled by the player.
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:CHARACTERS|>
|
||||
{% for character in characters %}
|
||||
### {{ character.name }}
|
||||
{% if max_tokens > 6000 -%}
|
||||
{{ character.sheet }}
|
||||
{% else -%}
|
||||
{{ character.filtered_sheet(['age', 'gender']) }}
|
||||
{{ query_memory("what is "+character.name+"'s personality?", as_question_answer=False) }}
|
||||
{% endif %}
|
||||
|
||||
{{ character.description }}
|
||||
{% endfor %}
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:TASK|>
|
||||
Generate the introductory text for the player as he starts this text based adventure game.
|
||||
|
||||
Use the premise to guide the text generation.
|
||||
|
||||
Start the player off in the beginning of the story and dont reveal too much information just yet.
|
||||
|
||||
The text must be short (200 words or less) and should be immersive.
|
||||
|
||||
Writh from a third person perspective and use the character names to refer to the characters.
|
||||
|
||||
The player, as Elmer, will see the text you generate when they first enter the game world.
|
||||
|
||||
The text should be immersive and should put the player into an actionable state. The ending of the text should be a prompt for the player's first action.
|
||||
<|CLOSE_SECTION|>
|
||||
{{ set_prepared_response('You') }}
|
||||
@@ -0,0 +1,36 @@
|
||||
<|SECTION:DESCRIPTION|>
|
||||
{{ scene.description }}
|
||||
|
||||
Elmer and Kaira are the only crew members of the Starlight Nomad, a small spaceship traveling through interstellar space.
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:CHARACTERS|>
|
||||
{% for character in characters %}
|
||||
### {{ character.name }}
|
||||
{% if max_tokens > 6000 -%}
|
||||
{{ character.sheet }}
|
||||
{% else -%}
|
||||
{{ character.filtered_sheet(['age', 'gender']) }}
|
||||
{{ query_memory("what is "+character.name+"'s personality?", as_question_answer=False) }}
|
||||
{% endif %}
|
||||
|
||||
{{ character.description }}
|
||||
{% endfor %}
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:TASK|>
|
||||
Your task is to write a scenario premise for a new infinity quest scenario. Think of it as a standalone episode that you are writing a preview for, setting the tone and main plot points.
|
||||
|
||||
This is for an open ended roleplaying game, so the scenario should be open ended as well.
|
||||
|
||||
Kaira and Elmer are the main characters. Elmer is controlled by the player.
|
||||
|
||||
Generate 2 paragraphs of text.
|
||||
|
||||
Use an informal and colloquial register with a conversational tone. Overall, the narrative is informal, conversational, natural, and spontaneous, with a sense of immediacy.
|
||||
|
||||
The scenario MUST BE contained to the Starlight Nomad spaceship. The spaceship is a small spaceship with a crew of 2.
|
||||
The scope of the story should be small and personal.
|
||||
|
||||
Thematic Tags: {{ thematic_tags }}
|
||||
Use the thematic tags to subtly guide your writing. The tags are not required to be used in the text, but should be used to guide your writing.
|
||||
<|CLOSE_SECTION|>
|
||||
{{ set_prepared_response('In this episode') }}
|
||||
@@ -0,0 +1,24 @@
|
||||
<|SECTION:PREMISE|>
|
||||
{{ scene.description }}
|
||||
|
||||
{{ premise }}
|
||||
Elmer and Kaira are the only crew members of the Starlight Nomad, a small spaceship traveling through interstellar space.
|
||||
|
||||
Kaira and Elmer are the main characters. Elmer is controlled by the player.
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:CHARACTERS|>
|
||||
{% for character in characters %}
|
||||
### {{ character.name }}
|
||||
{% if max_tokens > 6000 -%}
|
||||
{{ character.sheet }}
|
||||
{% else -%}
|
||||
{{ character.filtered_sheet(['age', 'gender']) }}
|
||||
{{ query_memory("what is "+character.name+"'s personality?", as_question_answer=False) }}
|
||||
{% endif %}
|
||||
|
||||
{{ character.description }}
|
||||
{% endfor %}
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:TASK|>
|
||||
Your task is to define one overarching, SIMPLE win codition for the provided infinity quest scenario. What does it mean to win this scenario? This should be a single sentence that can be evalulated as true or false.
|
||||
<|CLOSE_SECTION|>
|
||||
@@ -0,0 +1,42 @@
|
||||
{% set _ = debug("RUNNING GAME INSTRUCTS") -%}
|
||||
{% if not game_state.has_var('instr.premise') %}
|
||||
{# Generate scenario START #}
|
||||
|
||||
{%- set _ = emit_system("warning", "This is a dynamic scenario generation experiment for Infinity Quest. It will likely require a strong LLM to generate something coherent. GPT-4 or 34B+ if local. Temper your expectations.") -%}
|
||||
|
||||
{#- emit status update to the UX -#}
|
||||
{%- set _ = emit_status("busy", "Generating scenario ... [1/3]", as_scene_message=True) -%}
|
||||
|
||||
{#- thematic tags will be used to randomize generation -#}
|
||||
{%- set tags = thematic_generator.generate("color", "state_of_matter", "scifi_trope") -%}
|
||||
{# set tags = 'solid,meteorite,windy,theory' #}
|
||||
|
||||
{#- generate scenario premise -#}
|
||||
{%- set tmpl__scenario_premise = render_template('generate-scenario-premise', thematic_tags=tags) %}
|
||||
{%- set instr__premise = render_and_request(tmpl__scenario_premise) -%}
|
||||
|
||||
|
||||
{#- generate introductory text -#}
|
||||
{%- set _ = emit_status("busy", "Generating scenario ... [2/3]", as_scene_message=True) -%}
|
||||
{%- set tmpl__scenario_intro = render_template('generate-scenario-intro', premise=instr__premise) %}
|
||||
{%- set instr__intro = "*"+render_and_request(tmpl__scenario_intro)+"*" -%}
|
||||
|
||||
{#- generate win conditions -#}
|
||||
{%- set _ = emit_status("busy", "Generating scenario ... [3/3]", as_scene_message=True) -%}
|
||||
{%- set tmpl__win_conditions = render_template('generate-win-conditions', premise=instr__premise) %}
|
||||
{%- set instr__win_conditions = render_and_request(tmpl__win_conditions) -%}
|
||||
|
||||
{#- emit status update to the UX -#}
|
||||
{%- set status = emit_status("success", "Scenario ready.", as_scene_message=True) -%}
|
||||
|
||||
{# set gamestate variables #}
|
||||
{%- set _ = game_state.set_var("instr.premise", instr__premise, commit=True) -%}
|
||||
{%- set _ = game_state.set_var("instr.intro", instr__intro, commit=True) -%}
|
||||
{%- set _ = game_state.set_var("instr.win_conditions", instr__win_conditions, commit=True) -%}
|
||||
|
||||
{# set scene properties #}
|
||||
{%- set _ = scene.set_intro(instr__intro) -%}
|
||||
|
||||
{# Generate scenario END #}
|
||||
{% endif %}
|
||||
{# TODO: could do mid scene instructions here #}
|
||||
@@ -97,6 +97,7 @@
|
||||
"cover_image": null
|
||||
}
|
||||
],
|
||||
"immutable_save": true,
|
||||
"goal": null,
|
||||
"goals": [],
|
||||
"context": "an epic sci-fi adventure aimed at an adult audience.",
|
||||
|
||||
|
After Width: | Height: | Size: 1.7 MiB |
52
scenes/simulation-suite/simulation-suite.json
Normal file
@@ -0,0 +1,52 @@
|
||||
{
|
||||
"name": "Simulation Suite",
|
||||
"environment": "scene",
|
||||
"immutable_save": true,
|
||||
"restore_from": "simulation-suite.json",
|
||||
"experimental": true,
|
||||
"help": "Address the computer by starting your statements with 'Computer, ' followed by an instruction.\n\nExamples:\n'Computer, i would like to experience an adventure on a derelict space station'\n'Computer, add a horrific alien creature that is chasing me.'",
|
||||
"description": "",
|
||||
"intro": "*You have entered the simulation suite. No simulation is currently active and you are in a non-descript space with paneled walls surrounding you. The control panel next to you is pulsating with a green light, indicating readiness to receive a prompt to start the simulation.*",
|
||||
"archived_history": [],
|
||||
"history": [],
|
||||
"ts": "PT1S",
|
||||
"characters": [
|
||||
{
|
||||
"name": "You",
|
||||
"gender": "unknown",
|
||||
"color": "cornflowerblue",
|
||||
"base_attributes": {},
|
||||
"is_player": true
|
||||
}
|
||||
],
|
||||
"context": "a simulated experience",
|
||||
"game_state": {
|
||||
"ops":{
|
||||
"run_on_start": true,
|
||||
"always_direct": true
|
||||
},
|
||||
"variables": {}
|
||||
},
|
||||
"world_state": {
|
||||
"character_name_mappings": {
|
||||
"You": [
|
||||
"user",
|
||||
"player",
|
||||
"player character",
|
||||
"user character",
|
||||
"the user",
|
||||
"the player"
|
||||
]
|
||||
}
|
||||
},
|
||||
"assets": {
|
||||
"cover_image": "4b157dccac2ba71adb078a9d591f9900d6d62f3e86168a5e0e5e1e9faf6dc103",
|
||||
"assets": {
|
||||
"4b157dccac2ba71adb078a9d591f9900d6d62f3e86168a5e0e5e1e9faf6dc103": {
|
||||
"id": "4b157dccac2ba71adb078a9d591f9900d6d62f3e86168a5e0e5e1e9faf6dc103",
|
||||
"file_type": "png",
|
||||
"media_type": "image/png"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
118
scenes/simulation-suite/templates/computer.jinja2
Normal file
@@ -0,0 +1,118 @@
|
||||
<|SECTION:CONTEXT|>
|
||||
{% set scene_history=scene.context_history(budget=1024) %}
|
||||
{% for scene_context in scene_history -%}
|
||||
{{ loop.index }}. {{ scene_context }}
|
||||
{% endfor %}
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:FUNCTIONS|>
|
||||
The player has instructed the computer to alter the current simulation.
|
||||
|
||||
You have access to the following functions, you can call as many as you want to fulfill the player's requests.
|
||||
|
||||
You must at least call one of the following functions:
|
||||
|
||||
- change_environment
|
||||
- add_ai_character
|
||||
- change_ai_character
|
||||
- remove_ai_character
|
||||
- set_player_persona
|
||||
- set_player_name
|
||||
- end_simulation
|
||||
- answer_question
|
||||
|
||||
`add_ai_character` and `change_ai_character` are exclusive if they are targeting the same character.
|
||||
|
||||
Set the player persona at the beginning of a new simulation or if the player requests a change.
|
||||
|
||||
Only end the simulation if the player requests it explicitly.
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:EXAMPLES|>
|
||||
Request: Computer, I want to be on a mountain top
|
||||
```simulation-stack
|
||||
change_environment("mountain top")
|
||||
set_player_persona("mountain climber")
|
||||
set_player_name("Hank")
|
||||
```
|
||||
|
||||
Request: Computer, I want to be more muscular and taller
|
||||
```simulation-stack
|
||||
set_player_persona("make player more muscular and taller")
|
||||
```
|
||||
|
||||
Request: Computer, the building should be on fire
|
||||
```simulation-stack
|
||||
change_environment("building on fire")
|
||||
```
|
||||
|
||||
Request: Computer, a rocket hits the building and George is now injured
|
||||
```simulation-stack
|
||||
change_environment("building on fire")
|
||||
change_ai_character("George is injured")
|
||||
```
|
||||
|
||||
Request: Computer, I want to experience a rollercoaster ride with a friend
|
||||
```simulation-stack
|
||||
change_environment("theme park, riding a rollercoaster")
|
||||
set_player_persona("young female experiencing rollercoaster ride")
|
||||
set_player_name("Susanne")
|
||||
add_ai_character("a female friend of player named Sarah")
|
||||
```
|
||||
|
||||
Request: Computer, I want to experience the international space station
|
||||
```simulation-stack
|
||||
change_environment("international space station")
|
||||
set_player_persona("astronaut experiencing first trip to ISS")
|
||||
set_player_name("George")
|
||||
add_ai_character("astronaut named Henry")
|
||||
```
|
||||
|
||||
Request: Computer, remove the goblin and add an elven woman instead
|
||||
```simulation-stack
|
||||
remove_ai_character("goblin")
|
||||
add_ai_character("elven woman named Elune")
|
||||
```
|
||||
|
||||
Request: Computer, change the skiing instructor to be older.
|
||||
```simulation-stack
|
||||
change_ai_character("make skiing instructor older")
|
||||
```
|
||||
|
||||
Request: Computer, change my grandma to my grandpa
|
||||
```simulation-stack
|
||||
remove_ai_character("grandma")
|
||||
add_ai_character("grandpa named Steven")
|
||||
```
|
||||
|
||||
Request: Computer, remove the skiing instructor and add my friend instead.
|
||||
```simulation-stack
|
||||
remove_ai_character("skiing instructor")
|
||||
add_ai_character("player's friend named Tara")
|
||||
```
|
||||
|
||||
Request: Computer, replace the skiing instructor with my friend.
|
||||
```simulation-stack
|
||||
remove_ai_character("skiing instructor")
|
||||
add_ai_character("player's friend named Lisa")
|
||||
```
|
||||
|
||||
Request: Computer, I want to end the simulation
|
||||
```simulation-stack
|
||||
end_simulation("simulation ended")
|
||||
```
|
||||
|
||||
Request: Computer, shut down the simulation
|
||||
```simulation-stack
|
||||
end_simulation("simulation ended")
|
||||
```
|
||||
|
||||
Request: Computer, what do you know about the game of thrones?
|
||||
```simulation-stack
|
||||
answer_question("what do you know about the game of thrones?")
|
||||
```
|
||||
|
||||
<|CLOSE_SECTION|>
|
||||
<|SECTION:TASK|>
|
||||
Respond with the simulation stack for the following request:
|
||||
|
||||
Request: {{ player_instruction }}
|
||||
{{ bot_token }}```simulation-stack
|
||||
177
scenes/simulation-suite/templates/instructions.jinja2
Normal file
@@ -0,0 +1,177 @@
|
||||
{% set update_world_state = False %}
|
||||
{% set _ = debug("HOLODECK SIMULATION") -%}
|
||||
{% set player_character = scene.get_player_character() %}
|
||||
{% set player_message = scene.last_player_message() %}
|
||||
{% set last_processed = game_state.get_var('instr.last_processed', -1) %}
|
||||
{% set player_message_is_instruction = (player_message and player_message.raw.lower().startswith("computer") and not player_message.hidden) and not player_message.raw.lower().strip() == "computer" and not last_processed >= player_message.id %}
|
||||
{% set simulation_reset = False %}
|
||||
{% if not game_state.has_var('instr.simulation_stopped') %}
|
||||
{# simulation NOT started #}
|
||||
|
||||
{# get last player instruction #}
|
||||
{% if player_message_is_instruction %}
|
||||
{# player message exists #}
|
||||
|
||||
{#% set _ = agent_action("narrator", "action_to_narration", action_name="paraphrase", narration="The computer is processing the request, please wait a moment.", emit_message=True) %#}
|
||||
|
||||
{% set calls = render_and_request(render_template("computer", player_instruction=player_message.raw), dedupe_enabled=False) %}
|
||||
|
||||
{% set _ = debug("HOLODECK simulation calls", calls=calls ) %}
|
||||
{% set processed = make_list() %}
|
||||
|
||||
{% for call in calls.split("\n") %}
|
||||
{% set _ = debug("CALL", call=call, processed=processed) %}
|
||||
{% set inject = "The computer executes the function `"+call+"`" %}
|
||||
{% if call.strip().startswith('change_environment') %}
|
||||
{# change environment #}
|
||||
{% set _ = processed.append(call) %}
|
||||
|
||||
{% elif call.strip().startswith("answer_question") %}
|
||||
{# answert a query #}
|
||||
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="progress_story", narrative_direction="The computer calls the following function:\n"+call+"\nand answers the player's question.", emit_message=True) %}
|
||||
|
||||
|
||||
{% elif call.strip().startswith("set_player_persona") %}
|
||||
{# treansform player #}
|
||||
{% set _ = emit_status("busy", "Simulation suite altering user persona.", as_scene_message=True) %}
|
||||
|
||||
{% set character_attributes = agent_action("world_state", "extract_character_sheet", name=player_character.name, text=player_message.raw)%}
|
||||
|
||||
{% set _ = player_character.update(base_attributes=character_attributes) %}
|
||||
|
||||
{% set character_description = agent_action("creator", "determine_character_description", character=player_character) %}
|
||||
|
||||
{% set _ = player_character.update(description=character_description) %}
|
||||
|
||||
{% set _ = debug("HOLODECK transform player", attributes=character_attributes, description=character_description) %}
|
||||
{% set _ = processed.append(call) %}
|
||||
{% elif call.strip().startswith("set_player_name") %}
|
||||
{# change player name #}
|
||||
{% set _ = emit_status("busy", "Simulation suite adjusting user idenity.", as_scene_message=True) %}
|
||||
{% set character_name = agent_action("creator", "determine_character_name", character_name=inject+" - What is a fitting name for the player persona? Respond with the current name if it still fits.") %}
|
||||
|
||||
{% set _ = debug("HOLODECK player name", character_name=character_name) %}
|
||||
|
||||
{% if character_name != player_character.name %}
|
||||
{% set _ = processed.append(call) %}
|
||||
{% set _ = player_character.rename(character_name) %}
|
||||
{% endif %}
|
||||
{% elif call.strip().startswith("add_ai_character") %}
|
||||
{# add new npc #}
|
||||
|
||||
{% set _ = emit_status("busy", "Simulation suite adding character.", as_scene_message=True) %}
|
||||
{% set character_name = agent_action("creator", "determine_character_name", character_name=inject+" - what is the name of the character to be added to the scene? If no name can extracted from the text, extract a short descriptive name instead. Respond only with the name.") %}
|
||||
|
||||
{% set _ = emit_status("busy", "Simulation suite adding character: "+character_name, as_scene_message=True) %}
|
||||
{% set _ = debug("HOLODECK add npc", name=character_name)%}
|
||||
{% set npc = agent_action("director", "persist_character", name=character_name, content=player_message.raw )%}
|
||||
{% set _ = agent_action("world_state", "manager", action_name="add_detail_reinforcement", character_name=npc.name, question="Goal", instructions="Generate a goal for "+npc.name+", based on the user's chosen simulation", interval=25, run_immediately=True) %}
|
||||
{% set _ = debug("HOLODECK added npc", npc=npc) %}
|
||||
{% set _ = processed.append(call) %}
|
||||
{% set _ = agent_action("visual", "generate_character_portrait", character_name=npc.name) %}
|
||||
{% elif call.strip().startswith("remove_ai_character") %}
|
||||
{# remove npc #}
|
||||
|
||||
{% set _ = emit_status("busy", "Simulation suite removing character.", as_scene_message=True) %}
|
||||
{% set character_name = agent_action("creator", "determine_character_name", character_name=inject+" - what is the name of the character being removed?", allowed_names=scene.npc_character_names) %}
|
||||
|
||||
{% set npc = scene.get_character(character_name) %}
|
||||
|
||||
{% if npc %}
|
||||
{% set _ = debug("HOLODECK remove npc", npc=npc.name) %}
|
||||
{% set _ = agent_action("world_state", "manager", action_name="deactivate_character", character_name=npc.name) %}
|
||||
{% set _ = processed.append(call) %}
|
||||
{% endif %}
|
||||
{% elif call.strip().startswith("change_ai_character") %}
|
||||
{# change existing npc #}
|
||||
|
||||
{% set _ = emit_status("busy", "Simulation suite altering character.", as_scene_message=True) %}
|
||||
{% set character_name = agent_action("creator", "determine_character_name", character_name=inject+" - what is the name of the character receiving the changes (before the change)?", allowed_names=scene.npc_character_names) %}
|
||||
|
||||
{% set character_name_after = agent_action("creator", "determine_character_name", character_name=inject+" - what is the name of the character receiving the changes (after the changes)?") %}
|
||||
|
||||
{% set npc = scene.get_character(character_name) %}
|
||||
|
||||
{% if npc %}
|
||||
{% set _ = emit_status("busy", "Changing "+character_name+" -> "+character_name_after, as_scene_message=True) %}
|
||||
{% set _ = debug("HOLODECK transform npc", npc=npc) %}
|
||||
{% set character_attributes = agent_action("world_state", "extract_character_sheet", name=npc.name, alteration_instructions=player_message.raw)%}
|
||||
{% set _ = npc.update(base_attributes=character_attributes) %}
|
||||
{% set character_description = agent_action("creator", "determine_character_description", character=npc) %}
|
||||
{% set _ = npc.update(description=character_description) %}
|
||||
{% set _ = debug("HOLODECK transform npc", attributes=character_attributes, description=character_description) %}
|
||||
{% set _ = processed.append(call) %}
|
||||
{% if character_name_after != character_name %}
|
||||
{% set _ = npc.rename(character_name_after) %}
|
||||
{% endif %}
|
||||
{% endif %}
|
||||
{% elif call.strip().startswith("end_simulation") %}
|
||||
{# end simulation #}
|
||||
{% set explicit_command = query_text_eval("has the player explicitly asked to end the simulation?", player_message.raw) %}
|
||||
{% if explicit_command %}
|
||||
{% set _ = emit_status("busy", "Simulation suite ending current simulation.", as_scene_message=True) %}
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="progress_story", narrative_direction="The computer ends the simulation, disolving the environment and all artifical characters, erasing all memory of it and finally returning the player to the inactive simulation suite.List of artificial characters: "+(",".join(scene.npc_character_names))+". The player is also transformed back to their normal persona.", emit_message=True) %}
|
||||
{% set _ = scene.sync_restore() %}
|
||||
{% set _ = agent_action("world_state", "update_world_state", force=True) %}
|
||||
{% set simulation_reset = True %}
|
||||
{% endif %}
|
||||
{% elif "(" in call.strip() %}
|
||||
{# unknown function call, still add it to processed stack so it can be incoorporated in the narration #}
|
||||
{% set _ = processed.append(call) %}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
|
||||
{% if processed and not simulation_reset %}
|
||||
{% set _ = game_state.set_var("instr.has_issued_instructions", "yes", commit=False) %}
|
||||
{% set _ = emit_status("busy", "Simulation suite altering environment.", as_scene_message=True) %}
|
||||
{% set update_world_state = True %}
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="progress_story", narrative_direction="The computer calls the following functions:\n"+processed.join("\n")+"\nand the simulation adjusts the environment according to the user's wishes.\n\nWrite the narrative that describes the changes to the player in the context of the simulation starting up.", emit_message=True) %}
|
||||
{% endif %}
|
||||
|
||||
{% elif not game_state.has_var("instr.simulation_started") %}
|
||||
{# no player message yet, start of scenario #}
|
||||
{% set _ = emit_status("busy", "Simulation suite powering up.", as_scene_message=True) %}
|
||||
{% set _ = game_state.set_var("instr.simulation_started", "yes", commit=False) %}
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="progress_story", narrative_direction="Narrate the computer asking the user to state the nature of their desired simulation.", emit_message=False) %}
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="passthrough", narration="Please state your commands by addressing the computer by stating \"Computer,\" followed by an instruction.") %}
|
||||
|
||||
{# pin to make sure characters don't try to interact with the simulation #}
|
||||
{% set _ = agent_action("world_state", "manager", action_name="save_world_entry", entry_id="sim.quarantined", text="Characters in the simulation ARE NOT AWARE OF THE COMPUTER.", meta=make_dict(), pin=True) %}
|
||||
|
||||
{% set _ = emit_status("success", "Simulation suite ready", as_scene_message=True) %}
|
||||
{% endif %}
|
||||
|
||||
{% else %}
|
||||
{# simulation ongoing #}
|
||||
|
||||
{% endif %}
|
||||
|
||||
{% if update_world_state %}
|
||||
{% set _ = emit_status("busy", "Simulation suite updating world state.", as_scene_message=True) %}
|
||||
{% set _ = agent_action("world_state", "update_world_state", force=True) %}
|
||||
{% endif %}
|
||||
|
||||
{% if not scene.npc_character_names and not simulation_reset %}
|
||||
{# no characters in the scene, see if there are any to add #}
|
||||
{% set npcs = agent_action("director", "persist_characters_from_worldstate", exclude=["computer", "user", "player", "you"]) %}
|
||||
{% for npc in npcs %}
|
||||
{% set _ = agent_action("world_state", "manager", action_name="add_detail_reinforcement", character_name=npc.name, question="Goal", instructions="Generate a goal for the character, based on the user's chosen simulation", interval=25, run_immediately=True) %}
|
||||
{% endfor %}
|
||||
{% if npcs %}
|
||||
{% set _ = agent_action("world_state", "update_world_state", force=True) %}
|
||||
{% endif %}
|
||||
{% endif %}
|
||||
|
||||
{% if player_message_is_instruction %}
|
||||
{# hide player message to the computer, so its not included in the scene context #}
|
||||
{% set _ = player_message.hide() %}
|
||||
{% set _ = game_state.set_var("instr.last_processed", player_message.id, commit=False) %}
|
||||
{% set _ = emit_status("success", "Simulation suite processed instructions", as_scene_message=True) %}
|
||||
{% elif player_message and not game_state.has_var("instr.has_issued_instructions") %}
|
||||
{# simulation not started, but player message is not an instruction #}
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="paraphrase", narration="Instructions to the simulation computer are only process if the computer is addressed at the beginning of the instruction. Please state your commands by addressing the computer by stating \"Computer,\" followed by an instruction. For example ... \"Computer, i want to experience being on a derelict spaceship.\"", emit_message=True) %}
|
||||
{% elif player_message and not scene.npc_character_names %}
|
||||
{# simulation started, player message is NOT an instruction, but there are no npcs to interact with #}
|
||||
{% set _ = agent_action("narrator", "action_to_narration", action_name="progress_story", narrative_direction="The environment reacts to the player's actions. YOU MUST NOT ACT ON BEHALF OF THE PLAYER. YOU MUST NOT INTERACT WITH THE COMPUTER.", emit_message=True) %}
|
||||
{% endif %}
|
||||
@@ -2,4 +2,4 @@ from .agents import Agent
|
||||
from .client import TextGeneratorWebuiClient
|
||||
from .tale_mate import *
|
||||
|
||||
VERSION = "0.13.2"
|
||||
VERSION = "0.21.0"
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
from .base import Agent
|
||||
from .creator import CreatorAgent
|
||||
from .context import ContextAgent
|
||||
from .conversation import ConversationAgent
|
||||
from .creator import CreatorAgent
|
||||
from .director import DirectorAgent
|
||||
from .editor import EditorAgent
|
||||
from .memory import ChromaDBMemoryAgent, MemoryAgent
|
||||
from .narrator import NarratorAgent
|
||||
from .registry import AGENT_CLASSES, get_agent_class, register
|
||||
from .summarize import SummarizeAgent
|
||||
from .editor import EditorAgent
|
||||
from .world_state import WorldStateAgent
|
||||
from .tts import TTSAgent
|
||||
from .visual import VisualAgent
|
||||
from .world_state import WorldStateAgent
|
||||
|
||||
@@ -1,64 +1,96 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import re
|
||||
from abc import ABC
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
from blinker import signal
|
||||
|
||||
import talemate.instance as instance
|
||||
import talemate.util as util
|
||||
from talemate.emit import emit
|
||||
from talemate.events import GameLoopStartEvent
|
||||
import talemate.emit.async_signals
|
||||
import dataclasses
|
||||
import pydantic
|
||||
import structlog
|
||||
from blinker import signal
|
||||
|
||||
import talemate.emit.async_signals
|
||||
import talemate.instance as instance
|
||||
import talemate.util as util
|
||||
from talemate.agents.context import ActiveAgent
|
||||
from talemate.emit import emit
|
||||
from talemate.events import GameLoopStartEvent
|
||||
|
||||
__all__ = [
|
||||
"Agent",
|
||||
"AgentAction",
|
||||
"AgentActionConditional",
|
||||
"AgentActionConfig",
|
||||
"AgentDetail",
|
||||
"AgentEmission",
|
||||
"set_processing",
|
||||
]
|
||||
|
||||
log = structlog.get_logger("talemate.agents.base")
|
||||
|
||||
|
||||
class AgentActionConfig(pydantic.BaseModel):
|
||||
type: str
|
||||
label: str
|
||||
description: str = ""
|
||||
value: Union[int, float, str, bool]
|
||||
value: Union[int, float, str, bool, None] = None
|
||||
default_value: Union[int, float, str, bool] = None
|
||||
max: Union[int, float, None] = None
|
||||
min: Union[int, float, None] = None
|
||||
step: Union[int, float, None] = None
|
||||
scope: str = "global"
|
||||
choices: Union[list[dict[str, str]], None] = None
|
||||
note: Union[str, None] = None
|
||||
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
|
||||
class AgentActionConditional(pydantic.BaseModel):
|
||||
attribute: str
|
||||
value: Union[int, float, str, bool, None] = None
|
||||
|
||||
|
||||
class AgentAction(pydantic.BaseModel):
|
||||
enabled: bool = True
|
||||
label: str
|
||||
description: str = ""
|
||||
config: Union[dict[str, AgentActionConfig], None] = None
|
||||
|
||||
condition: Union[AgentActionConditional, None] = None
|
||||
|
||||
|
||||
class AgentDetail(pydantic.BaseModel):
|
||||
value: Union[str, None] = None
|
||||
description: Union[str, None] = None
|
||||
icon: Union[str, None] = None
|
||||
color: str = "grey"
|
||||
|
||||
|
||||
def set_processing(fn):
|
||||
"""
|
||||
decorator that emits the agent status as processing while the function
|
||||
is running.
|
||||
|
||||
|
||||
Done via a try - final block to ensure the status is reset even if
|
||||
the function fails.
|
||||
"""
|
||||
|
||||
|
||||
@wraps(fn)
|
||||
async def wrapper(self, *args, **kwargs):
|
||||
try:
|
||||
await self.emit_status(processing=True)
|
||||
return await fn(self, *args, **kwargs)
|
||||
finally:
|
||||
await self.emit_status(processing=False)
|
||||
|
||||
wrapper.__name__ = fn.__name__
|
||||
|
||||
with ActiveAgent(self, fn):
|
||||
try:
|
||||
await self.emit_status(processing=True)
|
||||
return await fn(self, *args, **kwargs)
|
||||
finally:
|
||||
try:
|
||||
await self.emit_status(processing=False)
|
||||
except RuntimeError as exc:
|
||||
# not sure why this happens
|
||||
# some concurrency error?
|
||||
log.error("error emitting agent status", exc=exc)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@@ -70,6 +102,11 @@ class Agent(ABC):
|
||||
agent_type = "agent"
|
||||
verbose_name = None
|
||||
set_processing = set_processing
|
||||
requires_llm_client = True
|
||||
auto_break_repetition = False
|
||||
websocket_handler = None
|
||||
essential = True
|
||||
ready_check_error = None
|
||||
|
||||
@property
|
||||
def agent_details(self):
|
||||
@@ -82,46 +119,51 @@ class Agent(ABC):
|
||||
def verbose_name(self):
|
||||
return self.agent_type.capitalize()
|
||||
|
||||
|
||||
|
||||
@property
|
||||
def ready(self):
|
||||
if not getattr(self.client, "enabled", True):
|
||||
return False
|
||||
|
||||
|
||||
if self.client and self.client.current_status in ["error", "warning"]:
|
||||
return False
|
||||
|
||||
|
||||
return self.client is not None
|
||||
|
||||
@property
|
||||
def status(self):
|
||||
if self.ready:
|
||||
if not self.enabled:
|
||||
return "disabled"
|
||||
return "idle" if getattr(self, "processing", 0) == 0 else "busy"
|
||||
else:
|
||||
if not self.enabled:
|
||||
return "disabled"
|
||||
|
||||
if not self.ready:
|
||||
return "uninitialized"
|
||||
|
||||
if getattr(self, "processing", 0) > 0:
|
||||
return "busy"
|
||||
|
||||
if getattr(self, "processing_bg", 0) > 0:
|
||||
return "busy_bg"
|
||||
|
||||
return "idle"
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
# by default, agents are enabled, an agent class that
|
||||
# is disableable should override this property
|
||||
return True
|
||||
|
||||
|
||||
@property
|
||||
def disable(self):
|
||||
# by default, agents are enabled, an agent class that
|
||||
# is disableable should override this property to
|
||||
# is disableable should override this property to
|
||||
# disable the agent
|
||||
pass
|
||||
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
# by default, agents do not have toggles to enable / disable
|
||||
# an agent class that is disableable should override this property
|
||||
return False
|
||||
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
# by default, agents are not experimental, an agent class that
|
||||
@@ -135,102 +177,173 @@ class Agent(ABC):
|
||||
"enabled": agent.enabled if agent else True,
|
||||
"has_toggle": agent.has_toggle if agent else False,
|
||||
"experimental": agent.experimental if agent else False,
|
||||
"requires_llm_client": cls.requires_llm_client,
|
||||
}
|
||||
actions = getattr(agent, "actions", None)
|
||||
|
||||
|
||||
if actions:
|
||||
config_options["actions"] = {k: v.model_dump() for k, v in actions.items()}
|
||||
else:
|
||||
config_options["actions"] = {}
|
||||
|
||||
|
||||
return config_options
|
||||
|
||||
def apply_config(self, *args, **kwargs):
|
||||
@property
|
||||
def meta(self):
|
||||
return {
|
||||
"essential": self.essential,
|
||||
}
|
||||
|
||||
async def _handle_ready_check(self, fut: asyncio.Future):
|
||||
callback_failure = getattr(self, "on_ready_check_failure", None)
|
||||
if fut.cancelled():
|
||||
if callback_failure:
|
||||
await callback_failure()
|
||||
return
|
||||
|
||||
if fut.exception():
|
||||
exc = fut.exception()
|
||||
self.ready_check_error = exc
|
||||
log.error("agent ready check error", agent=self.agent_type, exc=exc)
|
||||
if callback_failure:
|
||||
await callback_failure(exc)
|
||||
return
|
||||
|
||||
callback = getattr(self, "on_ready_check_success", None)
|
||||
if callback:
|
||||
await callback()
|
||||
|
||||
async def ready_check(self, task: asyncio.Task = None):
|
||||
self.ready_check_error = None
|
||||
if task:
|
||||
task.add_done_callback(
|
||||
lambda fut: asyncio.create_task(self._handle_ready_check(fut))
|
||||
)
|
||||
return
|
||||
return True
|
||||
|
||||
async def apply_config(self, *args, **kwargs):
|
||||
if self.has_toggle and "enabled" in kwargs:
|
||||
self.is_enabled = kwargs.get("enabled", False)
|
||||
|
||||
|
||||
if not getattr(self, "actions", None):
|
||||
return
|
||||
|
||||
|
||||
for action_key, action in self.actions.items():
|
||||
|
||||
if not kwargs.get("actions"):
|
||||
continue
|
||||
|
||||
action.enabled = kwargs.get("actions", {}).get(action_key, {}).get("enabled", False)
|
||||
|
||||
|
||||
action.enabled = (
|
||||
kwargs.get("actions", {}).get(action_key, {}).get("enabled", False)
|
||||
)
|
||||
|
||||
if not action.config:
|
||||
continue
|
||||
|
||||
|
||||
for config_key, config in action.config.items():
|
||||
try:
|
||||
config.value = kwargs.get("actions", {}).get(action_key, {}).get("config", {}).get(config_key, {}).get("value", config.value)
|
||||
config.value = (
|
||||
kwargs.get("actions", {})
|
||||
.get(action_key, {})
|
||||
.get("config", {})
|
||||
.get(config_key, {})
|
||||
.get("value", config.value)
|
||||
)
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
async def on_game_loop_start(self, event:GameLoopStartEvent):
|
||||
|
||||
|
||||
async def on_game_loop_start(self, event: GameLoopStartEvent):
|
||||
"""
|
||||
Finds all ActionConfigs that have a scope of "scene" and resets them to their default values
|
||||
"""
|
||||
|
||||
|
||||
if not getattr(self, "actions", None):
|
||||
return
|
||||
|
||||
|
||||
for _, action in self.actions.items():
|
||||
if not action.config:
|
||||
continue
|
||||
|
||||
|
||||
for _, config in action.config.items():
|
||||
if config.scope == "scene":
|
||||
# if default_value is None, just use the `type` of the current
|
||||
# if default_value is None, just use the `type` of the current
|
||||
# value
|
||||
if config.default_value is None:
|
||||
default_value = type(config.value)()
|
||||
else:
|
||||
default_value = config.default_value
|
||||
|
||||
log.debug("resetting config", config=config, default_value=default_value)
|
||||
|
||||
log.debug(
|
||||
"resetting config", config=config, default_value=default_value
|
||||
)
|
||||
config.value = default_value
|
||||
|
||||
|
||||
await self.emit_status()
|
||||
|
||||
|
||||
async def emit_status(self, processing: bool = None):
|
||||
|
||||
# should keep a count of processing requests, and when the
|
||||
# number is 0 status is "idle", if the number is greater than 0
|
||||
# status is "busy"
|
||||
#
|
||||
# increase / decrease based on value of `processing`
|
||||
|
||||
|
||||
if getattr(self, "processing", None) is None:
|
||||
self.processing = 0
|
||||
|
||||
if not processing:
|
||||
|
||||
if processing is False:
|
||||
self.processing -= 1
|
||||
self.processing = max(0, self.processing)
|
||||
else:
|
||||
elif processing is True:
|
||||
self.processing += 1
|
||||
|
||||
status = "busy" if self.processing > 0 else "idle"
|
||||
if not self.enabled:
|
||||
status = "disabled"
|
||||
|
||||
|
||||
emit(
|
||||
"agent_status",
|
||||
message=self.verbose_name or "",
|
||||
id=self.agent_type,
|
||||
status=status,
|
||||
status=self.status,
|
||||
details=self.agent_details,
|
||||
meta=self.meta,
|
||||
data=self.config_options(agent=self),
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def _handle_background_processing(self, fut: asyncio.Future):
|
||||
try:
|
||||
if fut.cancelled():
|
||||
return
|
||||
|
||||
if fut.exception():
|
||||
log.error(
|
||||
"background processing error",
|
||||
agent=self.agent_type,
|
||||
exc=fut.exception(),
|
||||
)
|
||||
await self.emit_status()
|
||||
return
|
||||
|
||||
log.info("background processing done", agent=self.agent_type)
|
||||
finally:
|
||||
self.processing_bg -= 1
|
||||
await self.emit_status()
|
||||
|
||||
async def set_background_processing(self, task: asyncio.Task):
|
||||
log.info("set_background_processing", agent=self.agent_type)
|
||||
if not hasattr(self, "processing_bg"):
|
||||
self.processing_bg = 0
|
||||
|
||||
self.processing_bg += 1
|
||||
|
||||
await self.emit_status()
|
||||
task.add_done_callback(
|
||||
lambda fut: asyncio.create_task(self._handle_background_processing(fut))
|
||||
)
|
||||
|
||||
def connect(self, scene):
|
||||
self.scene = scene
|
||||
talemate.emit.async_signals.get("game_loop_start").connect(self.on_game_loop_start)
|
||||
|
||||
talemate.emit.async_signals.get("game_loop_start").connect(
|
||||
self.on_game_loop_start
|
||||
)
|
||||
|
||||
def clean_result(self, result):
|
||||
if "#" in result:
|
||||
@@ -276,6 +389,27 @@ class Agent(ABC):
|
||||
current_memory_context.append(memory)
|
||||
return current_memory_context
|
||||
|
||||
# LLM client related methods. These are called during or after the client
|
||||
# sends the prompt to the API.
|
||||
|
||||
def inject_prompt_paramters(
|
||||
self, prompt_param: dict, kind: str, agent_function_name: str
|
||||
):
|
||||
"""
|
||||
Injects prompt parameters before the client sends off the prompt
|
||||
Override as needed.
|
||||
"""
|
||||
pass
|
||||
|
||||
def allow_repetition_break(
|
||||
self, kind: str, agent_function_name: str, auto: bool = False
|
||||
):
|
||||
"""
|
||||
Returns True if repetition breaking is allowed, False otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AgentEmission:
|
||||
agent: Agent
|
||||
agent: Agent
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
"""
|
||||
Code has been moved.
|
||||
"""
|
||||
"""
|
||||
|
||||
@@ -1,54 +1,46 @@
|
||||
from .base import Agent
|
||||
from .registry import register
|
||||
import contextvars
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
|
||||
import pydantic
|
||||
|
||||
__all__ = [
|
||||
"active_agent",
|
||||
]
|
||||
|
||||
active_agent = contextvars.ContextVar("active_agent", default=None)
|
||||
|
||||
|
||||
@register
|
||||
class ContextAgent(Agent):
|
||||
"""
|
||||
Agent that helps retrieve context for the continuation
|
||||
of dialogue.
|
||||
"""
|
||||
class ActiveAgentContext(pydantic.BaseModel):
|
||||
agent: object
|
||||
fn: Callable
|
||||
agent_stack: list = pydantic.Field(default_factory=list)
|
||||
|
||||
agent_type = "context"
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
def __init__(self, client, **kwargs):
|
||||
self.client = client
|
||||
@property
|
||||
def action(self):
|
||||
return self.fn.__name__
|
||||
|
||||
def determine_questions(self, scene_text):
|
||||
prompt = [
|
||||
"You are tasked to continue the following dialogue in a roleplaying session, but before you can do so you can ask three questions for extra context."
|
||||
"",
|
||||
"What are the questions you would ask?",
|
||||
"",
|
||||
"Known context and dialogue:" "",
|
||||
scene_text,
|
||||
"",
|
||||
"Questions:",
|
||||
"",
|
||||
]
|
||||
def __str__(self):
|
||||
return f"{self.agent.verbose_name}.{self.action}"
|
||||
|
||||
prompt = "\n".join(prompt)
|
||||
|
||||
questions = self.client.send_prompt(prompt, kind="question")
|
||||
class ActiveAgent:
|
||||
def __init__(self, agent, fn):
|
||||
self.agent = ActiveAgentContext(agent=agent, fn=fn)
|
||||
|
||||
questions = self.clean_result(questions)
|
||||
def __enter__(self):
|
||||
|
||||
return questions.split("\n")
|
||||
previous_agent = active_agent.get()
|
||||
|
||||
def get_answer(self, question, context):
|
||||
prompt = [
|
||||
"Read the context and answer the question:",
|
||||
"",
|
||||
"Context:",
|
||||
"",
|
||||
context,
|
||||
"",
|
||||
f"Question: {question}",
|
||||
"Answer:",
|
||||
]
|
||||
if previous_agent:
|
||||
self.agent.agent_stack = previous_agent.agent_stack + [str(self.agent)]
|
||||
else:
|
||||
self.agent.agent_stack = [str(self.agent)]
|
||||
|
||||
prompt = "\n".join(prompt)
|
||||
self.token = active_agent.set(self.agent)
|
||||
|
||||
answer = self.client.send_prompt(prompt, kind="answer")
|
||||
answer = self.clean_result(answer)
|
||||
return answer
|
||||
def __exit__(self, *args, **kwargs):
|
||||
active_agent.reset(self.token)
|
||||
return False
|
||||
|
||||
@@ -1,40 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import re
|
||||
import random
|
||||
import re
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Optional, Union
|
||||
|
||||
import structlog
|
||||
|
||||
import talemate.client as client
|
||||
import talemate.emit.async_signals
|
||||
import talemate.instance as instance
|
||||
import talemate.util as util
|
||||
import structlog
|
||||
from talemate.client.context import (
|
||||
client_context_attribute,
|
||||
set_client_context_attribute,
|
||||
set_conversation_context_attribute,
|
||||
)
|
||||
from talemate.emit import emit
|
||||
import talemate.emit.async_signals
|
||||
from talemate.scene_message import CharacterMessage, DirectorMessage
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.events import GameLoopEvent
|
||||
from talemate.client.context import set_conversation_context_attribute, client_context_attribute, set_client_context_attribute
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import CharacterMessage, DirectorMessage
|
||||
|
||||
from .base import Agent, AgentEmission, set_processing, AgentAction, AgentActionConfig
|
||||
from .base import Agent, AgentAction, AgentActionConfig, AgentEmission, set_processing
|
||||
from .registry import register
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate.tale_mate import Character, Scene, Actor
|
||||
|
||||
from talemate.tale_mate import Actor, Character, Scene
|
||||
|
||||
log = structlog.get_logger("talemate.agents.conversation")
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ConversationAgentEmission(AgentEmission):
|
||||
actor: Actor
|
||||
character: Character
|
||||
generation: list[str]
|
||||
|
||||
|
||||
|
||||
talemate.emit.async_signals.register(
|
||||
"agent.conversation.before_generate",
|
||||
"agent.conversation.generated"
|
||||
"agent.conversation.before_generate", "agent.conversation.generated"
|
||||
)
|
||||
|
||||
|
||||
@register()
|
||||
class ConversationAgent(Agent):
|
||||
"""
|
||||
@@ -45,7 +53,7 @@ class ConversationAgent(Agent):
|
||||
|
||||
agent_type = "conversation"
|
||||
verbose_name = "Conversation"
|
||||
|
||||
|
||||
min_dialogue_length = 75
|
||||
|
||||
def __init__(
|
||||
@@ -60,55 +68,69 @@ class ConversationAgent(Agent):
|
||||
self.logging_enabled = logging_enabled
|
||||
self.logging_date = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
self.current_memory_context = None
|
||||
|
||||
|
||||
# several agents extend this class, but we only want to initialize
|
||||
# these actions for the conversation agent
|
||||
|
||||
|
||||
if self.agent_type != "conversation":
|
||||
return
|
||||
|
||||
|
||||
self.actions = {
|
||||
"generation_override": AgentAction(
|
||||
enabled = True,
|
||||
label = "Generation Override",
|
||||
description = "Override generation parameters",
|
||||
config = {
|
||||
enabled=True,
|
||||
label="Generation Settings",
|
||||
config={
|
||||
"format": AgentActionConfig(
|
||||
type="text",
|
||||
label="Format",
|
||||
description="The format of the dialogue, as seen by the AI.",
|
||||
choices=[
|
||||
{"label": "Movie Script", "value": "movie_script"},
|
||||
{"label": "Chat (legacy)", "value": "chat"},
|
||||
],
|
||||
value="chat",
|
||||
),
|
||||
"length": AgentActionConfig(
|
||||
type="number",
|
||||
label="Generation Length (tokens)",
|
||||
description="Maximum number of tokens to generate for a conversation response.",
|
||||
value=96,
|
||||
value=128,
|
||||
min=32,
|
||||
max=512,
|
||||
step=32,
|
||||
),#
|
||||
), #
|
||||
"instructions": AgentActionConfig(
|
||||
type="text",
|
||||
label="Instructions",
|
||||
value="1-3 sentences.",
|
||||
value="Write 1-3 sentences. Never wax poetic.",
|
||||
description="Extra instructions to give the AI for dialog generatrion.",
|
||||
),
|
||||
"jiggle": AgentActionConfig(
|
||||
type="number",
|
||||
label="Jiggle",
|
||||
label="Jiggle (Increased Randomness)",
|
||||
description="If > 0.0 will cause certain generation parameters to have a slight random offset applied to them. The bigger the number, the higher the potential offset.",
|
||||
value=0.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.1,
|
||||
),
|
||||
}
|
||||
},
|
||||
),
|
||||
"auto_break_repetition": AgentAction(
|
||||
enabled=True,
|
||||
label="Auto Break Repetition",
|
||||
description="Will attempt to automatically break AI repetition.",
|
||||
),
|
||||
"natural_flow": AgentAction(
|
||||
enabled = True,
|
||||
label = "Natural Flow",
|
||||
description = "Will attempt to generate a more natural flow of conversation between multiple characters.",
|
||||
config = {
|
||||
enabled=True,
|
||||
label="Natural Flow",
|
||||
description="Will attempt to generate a more natural flow of conversation between multiple characters.",
|
||||
config={
|
||||
"max_auto_turns": AgentActionConfig(
|
||||
type="number",
|
||||
label="Max. Auto Turns",
|
||||
description="The maximum number of turns the AI is allowed to generate before it stops and waits for the player to respond.",
|
||||
value=4,
|
||||
value=4,
|
||||
min=1,
|
||||
max=100,
|
||||
step=1,
|
||||
@@ -117,67 +139,84 @@ class ConversationAgent(Agent):
|
||||
type="number",
|
||||
label="Max. Idle Turns",
|
||||
description="The maximum number of turns a character can go without speaking before they are considered overdue to speak.",
|
||||
value=8,
|
||||
value=8,
|
||||
min=1,
|
||||
max=100,
|
||||
step=1,
|
||||
),
|
||||
}
|
||||
},
|
||||
),
|
||||
"use_long_term_memory": AgentAction(
|
||||
enabled = True,
|
||||
label = "Long Term Memory",
|
||||
description = "Will augment the conversation prompt with long term memory.",
|
||||
config = {
|
||||
"ai_selected": AgentActionConfig(
|
||||
type="bool",
|
||||
label="AI Selected",
|
||||
description="If enabled, the AI will select the long term memory to use. (will increase how long it takes to generate a response)",
|
||||
value=False,
|
||||
enabled=True,
|
||||
label="Long Term Memory",
|
||||
description="Will augment the conversation prompt with long term memory.",
|
||||
config={
|
||||
"retrieval_method": AgentActionConfig(
|
||||
type="text",
|
||||
label="Context Retrieval Method",
|
||||
description="How relevant context is retrieved from the long term memory.",
|
||||
value="direct",
|
||||
choices=[
|
||||
{
|
||||
"label": "Context queries based on recent dialogue (fast)",
|
||||
"value": "direct",
|
||||
},
|
||||
{
|
||||
"label": "Context queries generated by AI",
|
||||
"value": "queries",
|
||||
},
|
||||
{
|
||||
"label": "AI compiled question and answers (slow)",
|
||||
"value": "questions",
|
||||
},
|
||||
],
|
||||
),
|
||||
}
|
||||
),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
@property
|
||||
def conversation_format(self):
|
||||
if self.actions["generation_override"].enabled:
|
||||
return self.actions["generation_override"].config["format"].value
|
||||
return "movie_script"
|
||||
|
||||
def connect(self, scene):
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("game_loop").connect(self.on_game_loop)
|
||||
|
||||
def last_spoken(self):
|
||||
|
||||
"""
|
||||
Returns the last time each character spoke
|
||||
"""
|
||||
|
||||
|
||||
last_turn = {}
|
||||
turns = 0
|
||||
character_names = self.scene.character_names
|
||||
max_idle_turns = self.actions["natural_flow"].config["max_idle_turns"].value
|
||||
|
||||
|
||||
for idx in range(len(self.scene.history) - 1, -1, -1):
|
||||
|
||||
if isinstance(self.scene.history[idx], CharacterMessage):
|
||||
|
||||
if turns >= max_idle_turns:
|
||||
break
|
||||
|
||||
|
||||
character = self.scene.history[idx].character_name
|
||||
|
||||
|
||||
if character in character_names:
|
||||
last_turn[character] = turns
|
||||
character_names.remove(character)
|
||||
|
||||
|
||||
if not character_names:
|
||||
break
|
||||
|
||||
|
||||
turns += 1
|
||||
|
||||
|
||||
if character_names and turns >= max_idle_turns:
|
||||
for character in character_names:
|
||||
last_turn[character] = max_idle_turns
|
||||
last_turn[character] = max_idle_turns
|
||||
|
||||
return last_turn
|
||||
|
||||
|
||||
def repeated_speaker(self):
|
||||
"""
|
||||
Counts the amount of times the most recent speaker has spoken in a row
|
||||
@@ -193,109 +232,164 @@ class ConversationAgent(Agent):
|
||||
else:
|
||||
break
|
||||
return count
|
||||
|
||||
async def on_game_loop(self, event:GameLoopEvent):
|
||||
|
||||
async def on_game_loop(self, event: GameLoopEvent):
|
||||
await self.apply_natural_flow()
|
||||
|
||||
async def apply_natural_flow(self):
|
||||
|
||||
async def apply_natural_flow(self, force: bool = False, npcs_only: bool = False):
|
||||
"""
|
||||
If the natural flow action is enabled, this will attempt to determine
|
||||
the ideal character to talk next.
|
||||
|
||||
|
||||
This will let the AI pick a character to talk to, but if the AI can't figure
|
||||
it out it will apply rules based on max_idle_turns and max_auto_turns.
|
||||
|
||||
|
||||
If all fails it will just pick a random character.
|
||||
|
||||
|
||||
Repetition is also taken into account, so if a character has spoken twice in a row
|
||||
they will not be picked again until someone else has spoken.
|
||||
"""
|
||||
|
||||
|
||||
scene = self.scene
|
||||
|
||||
if not scene.auto_progress and not force:
|
||||
# we only apply natural flow if auto_progress is enabled
|
||||
return
|
||||
|
||||
if self.actions["natural_flow"].enabled and len(scene.character_names) > 2:
|
||||
|
||||
# last time each character spoke (turns ago)
|
||||
max_idle_turns = self.actions["natural_flow"].config["max_idle_turns"].value
|
||||
max_auto_turns = self.actions["natural_flow"].config["max_auto_turns"].value
|
||||
last_turn = self.last_spoken()
|
||||
last_turn_player = last_turn.get(scene.get_player_character().name, 0)
|
||||
|
||||
if last_turn_player >= max_auto_turns:
|
||||
player_name = scene.get_player_character().name
|
||||
last_turn_player = last_turn.get(player_name, 0)
|
||||
|
||||
if last_turn_player >= max_auto_turns and not npcs_only:
|
||||
self.scene.next_actor = scene.get_player_character().name
|
||||
log.debug("conversation_agent.natural_flow", next_actor="player", overdue=True, player_character=scene.get_player_character().name)
|
||||
log.debug(
|
||||
"conversation_agent.natural_flow",
|
||||
next_actor="player",
|
||||
overdue=True,
|
||||
player_character=scene.get_player_character().name,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
log.debug("conversation_agent.natural_flow", last_turn=last_turn)
|
||||
|
||||
|
||||
# determine random character to talk, this will be the fallback in case
|
||||
# the AI can't figure out who should talk next
|
||||
|
||||
|
||||
if scene.prev_actor:
|
||||
|
||||
# we dont want to talk to the same person twice in a row
|
||||
character_names = scene.character_names
|
||||
character_names.remove(scene.prev_actor)
|
||||
|
||||
if npcs_only:
|
||||
character_names = [c for c in character_names if c != player_name]
|
||||
|
||||
random_character_name = random.choice(character_names)
|
||||
else:
|
||||
character_names = scene.character_names
|
||||
character_names = scene.character_names
|
||||
# no one has talked yet, so we just pick a random character
|
||||
|
||||
|
||||
if npcs_only:
|
||||
character_names = [c for c in character_names if c != player_name]
|
||||
|
||||
random_character_name = random.choice(scene.character_names)
|
||||
|
||||
overdue_characters = [character for character, turn in last_turn.items() if turn >= max_idle_turns]
|
||||
|
||||
|
||||
overdue_characters = [
|
||||
character
|
||||
for character, turn in last_turn.items()
|
||||
if turn >= max_idle_turns
|
||||
]
|
||||
|
||||
if npcs_only:
|
||||
overdue_characters = [c for c in overdue_characters if c != player_name]
|
||||
|
||||
if overdue_characters and self.scene.history:
|
||||
# Pick a random character from the overdue characters
|
||||
scene.next_actor = random.choice(overdue_characters)
|
||||
elif scene.history:
|
||||
scene.next_actor = None
|
||||
|
||||
|
||||
# AI will attempt to figure out who should talk next
|
||||
next_actor = await self.select_talking_actor(character_names)
|
||||
next_actor = next_actor.strip().strip('"').strip(".")
|
||||
|
||||
|
||||
for character_name in scene.character_names:
|
||||
if next_actor.lower() in character_name.lower() or character_name.lower() in next_actor.lower():
|
||||
if (
|
||||
next_actor.lower() in character_name.lower()
|
||||
or character_name.lower() in next_actor.lower()
|
||||
):
|
||||
scene.next_actor = character_name
|
||||
break
|
||||
|
||||
|
||||
if not scene.next_actor:
|
||||
# AI couldn't figure out who should talk next, so we just pick a random character
|
||||
log.debug("conversation_agent.natural_flow", next_actor="random", random_character_name=random_character_name)
|
||||
log.debug(
|
||||
"conversation_agent.natural_flow",
|
||||
next_actor="random",
|
||||
random_character_name=random_character_name,
|
||||
)
|
||||
scene.next_actor = random_character_name
|
||||
else:
|
||||
log.debug("conversation_agent.natural_flow", next_actor="picked", ai_next_actor=scene.next_actor)
|
||||
log.debug(
|
||||
"conversation_agent.natural_flow",
|
||||
next_actor="picked",
|
||||
ai_next_actor=scene.next_actor,
|
||||
)
|
||||
else:
|
||||
# always start with main character (TODO: configurable?)
|
||||
player_character = scene.get_player_character()
|
||||
log.debug("conversation_agent.natural_flow", next_actor="main_character", main_character=player_character)
|
||||
scene.next_actor = player_character.name if player_character else random_character_name
|
||||
|
||||
scene.log.debug("conversation_agent.natural_flow", next_actor=scene.next_actor)
|
||||
log.debug(
|
||||
"conversation_agent.natural_flow",
|
||||
next_actor="main_character",
|
||||
main_character=player_character,
|
||||
)
|
||||
scene.next_actor = (
|
||||
player_character.name if player_character else random_character_name
|
||||
)
|
||||
|
||||
scene.log.debug(
|
||||
"conversation_agent.natural_flow", next_actor=scene.next_actor
|
||||
)
|
||||
|
||||
|
||||
# same character cannot go thrice in a row, if this is happening, pick a random character that
|
||||
# isnt the same as the last character
|
||||
|
||||
if self.repeated_speaker() >= 2 and self.scene.prev_actor == self.scene.next_actor:
|
||||
scene.next_actor = random.choice([c for c in scene.character_names if c != scene.prev_actor])
|
||||
scene.log.debug("conversation_agent.natural_flow", next_actor="random (repeated safeguard)", random_character_name=scene.next_actor)
|
||||
|
||||
|
||||
if (
|
||||
self.repeated_speaker() >= 2
|
||||
and self.scene.prev_actor == self.scene.next_actor
|
||||
):
|
||||
scene.next_actor = random.choice(
|
||||
[c for c in scene.character_names if c != scene.prev_actor]
|
||||
)
|
||||
scene.log.debug(
|
||||
"conversation_agent.natural_flow",
|
||||
next_actor="random (repeated safeguard)",
|
||||
random_character_name=scene.next_actor,
|
||||
)
|
||||
|
||||
else:
|
||||
scene.next_actor = None
|
||||
|
||||
|
||||
@set_processing
|
||||
async def select_talking_actor(self, character_names: list[str]=None):
|
||||
result = await Prompt.request("conversation.select-talking-actor", self.client, "conversation_select_talking_actor", vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"character_names": character_names or self.scene.character_names,
|
||||
"character_names_formatted": ", ".join(character_names or self.scene.character_names),
|
||||
})
|
||||
|
||||
async def select_talking_actor(self, character_names: list[str] = None):
|
||||
result = await Prompt.request(
|
||||
"conversation.select-talking-actor",
|
||||
self.client,
|
||||
"conversation_select_talking_actor",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"character_names": character_names or self.scene.character_names,
|
||||
"character_names_formatted": ", ".join(
|
||||
character_names or self.scene.character_names
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def build_prompt_default(
|
||||
self,
|
||||
@@ -309,19 +403,17 @@ class ConversationAgent(Agent):
|
||||
# we subtract 200 to account for the response
|
||||
|
||||
scene = character.actor.scene
|
||||
|
||||
|
||||
total_token_budget = self.client.max_token_length - 200
|
||||
scene_and_dialogue_budget = total_token_budget - 500
|
||||
long_term_memory_budget = min(int(total_token_budget * 0.05), 200)
|
||||
|
||||
|
||||
scene_and_dialogue = scene.context_history(
|
||||
budget=scene_and_dialogue_budget,
|
||||
min_dialogue=25,
|
||||
budget=scene_and_dialogue_budget,
|
||||
keep_director=True,
|
||||
sections=False,
|
||||
insert_bot_token=10
|
||||
)
|
||||
|
||||
|
||||
memory = await self.build_prompt_default_memory(character)
|
||||
|
||||
main_character = scene.main_character.character
|
||||
@@ -336,39 +428,39 @@ class ConversationAgent(Agent):
|
||||
)
|
||||
else:
|
||||
formatted_names = character_names[0] if character_names else ""
|
||||
|
||||
# if there is more than 10 lines in scene_and_dialogue insert
|
||||
# a <|BOT|> token at -10, otherwise insert it at 0
|
||||
|
||||
try:
|
||||
director_message = isinstance(scene_and_dialogue[-1], DirectorMessage)
|
||||
except IndexError:
|
||||
director_message = False
|
||||
|
||||
|
||||
extra_instructions = ""
|
||||
if self.actions["generation_override"].enabled:
|
||||
extra_instructions = self.actions["generation_override"].config["instructions"].value
|
||||
extra_instructions = (
|
||||
self.actions["generation_override"].config["instructions"].value
|
||||
)
|
||||
|
||||
prompt = Prompt.get(
|
||||
"conversation.dialogue",
|
||||
vars={
|
||||
"scene": scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"scene_and_dialogue_budget": scene_and_dialogue_budget,
|
||||
"scene_and_dialogue": scene_and_dialogue,
|
||||
"memory": memory,
|
||||
"characters": list(scene.get_characters()),
|
||||
"main_character": main_character,
|
||||
"formatted_names": formatted_names,
|
||||
"talking_character": character,
|
||||
"partial_message": char_message,
|
||||
"director_message": director_message,
|
||||
"extra_instructions": extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
prompt = Prompt.get("conversation.dialogue", vars={
|
||||
"scene": scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"scene_and_dialogue_budget": scene_and_dialogue_budget,
|
||||
"scene_and_dialogue": scene_and_dialogue,
|
||||
"memory": memory,
|
||||
"characters": list(scene.get_characters()),
|
||||
"main_character": main_character,
|
||||
"formatted_names": formatted_names,
|
||||
"talking_character": character,
|
||||
"partial_message": char_message,
|
||||
"director_message": director_message,
|
||||
"extra_instructions": extra_instructions,
|
||||
})
|
||||
|
||||
return str(prompt)
|
||||
|
||||
async def build_prompt_default_memory(
|
||||
self, character: Character
|
||||
):
|
||||
|
||||
async def build_prompt_default_memory(self, character: Character):
|
||||
"""
|
||||
Builds long term memory for the conversation prompt
|
||||
|
||||
@@ -383,31 +475,56 @@ class ConversationAgent(Agent):
|
||||
if not self.actions["use_long_term_memory"].enabled:
|
||||
return []
|
||||
|
||||
|
||||
if self.current_memory_context:
|
||||
return self.current_memory_context
|
||||
|
||||
self.current_memory_context = ""
|
||||
|
||||
retrieval_method = (
|
||||
self.actions["use_long_term_memory"].config["retrieval_method"].value
|
||||
)
|
||||
|
||||
if self.actions["use_long_term_memory"].config["ai_selected"].value:
|
||||
history = self.scene.context_history(min_dialogue=3, max_dialogue=15, keep_director=False, sections=False, add_archieved_history=False)
|
||||
text = "\n".join(history)
|
||||
if retrieval_method != "direct":
|
||||
world_state = instance.get_agent("world_state")
|
||||
log.debug("conversation_agent.build_prompt_default_memory", direct=False)
|
||||
self.current_memory_context = await world_state.analyze_text_and_extract_context(
|
||||
text, f"continue the conversation as {character.name}"
|
||||
history = self.scene.context_history(
|
||||
min_dialogue=3,
|
||||
max_dialogue=15,
|
||||
keep_director=False,
|
||||
sections=False,
|
||||
add_archieved_history=False,
|
||||
)
|
||||
text = "\n".join(history)
|
||||
log.debug(
|
||||
"conversation_agent.build_prompt_default_memory",
|
||||
direct=False,
|
||||
version=retrieval_method,
|
||||
)
|
||||
|
||||
if retrieval_method == "questions":
|
||||
self.current_memory_context = (
|
||||
await world_state.analyze_text_and_extract_context(
|
||||
text, f"continue the conversation as {character.name}"
|
||||
)
|
||||
).split("\n")
|
||||
elif retrieval_method == "queries":
|
||||
self.current_memory_context = (
|
||||
await world_state.analyze_text_and_extract_context_via_queries(
|
||||
text, f"continue the conversation as {character.name}"
|
||||
)
|
||||
)
|
||||
|
||||
else:
|
||||
history = self.scene.context_history(min_dialogue=3, max_dialogue=3, keep_director=False, sections=False, add_archieved_history=False)
|
||||
log.debug("conversation_agent.build_prompt_default_memory", history=history, direct=True)
|
||||
history = list(map(str, self.scene.collect_messages(max_iterations=3)))
|
||||
log.debug(
|
||||
"conversation_agent.build_prompt_default_memory",
|
||||
history=history,
|
||||
direct=True,
|
||||
)
|
||||
memory = instance.get_agent("memory")
|
||||
|
||||
|
||||
context = await memory.multi_query(history, max_tokens=500, iterate=5)
|
||||
|
||||
self.current_memory_context = "\n".join(context)
|
||||
|
||||
|
||||
self.current_memory_context = context
|
||||
|
||||
return self.current_memory_context
|
||||
|
||||
async def build_prompt(self, character, char_message: str = ""):
|
||||
@@ -416,10 +533,9 @@ class ConversationAgent(Agent):
|
||||
return await fn(character, char_message=char_message)
|
||||
|
||||
def clean_result(self, result, character):
|
||||
|
||||
if "#" in result:
|
||||
result = result.split("#")[0]
|
||||
|
||||
|
||||
result = result.replace(" :", ":")
|
||||
result = result.replace("[", "*").replace("]", "*")
|
||||
result = result.replace("(", "*").replace(")", "*")
|
||||
@@ -430,31 +546,38 @@ class ConversationAgent(Agent):
|
||||
def set_generation_overrides(self):
|
||||
if not self.actions["generation_override"].enabled:
|
||||
return
|
||||
|
||||
set_conversation_context_attribute("length", self.actions["generation_override"].config["length"].value)
|
||||
|
||||
|
||||
set_conversation_context_attribute(
|
||||
"length", self.actions["generation_override"].config["length"].value
|
||||
)
|
||||
|
||||
if self.actions["generation_override"].config["jiggle"].value > 0.0:
|
||||
nuke_repetition = client_context_attribute("nuke_repetition")
|
||||
if nuke_repetition == 0.0:
|
||||
# we only apply the agent override if some other mechanism isn't already
|
||||
# setting the nuke_repetition value
|
||||
nuke_repetition = self.actions["generation_override"].config["jiggle"].value
|
||||
nuke_repetition = (
|
||||
self.actions["generation_override"].config["jiggle"].value
|
||||
)
|
||||
set_client_context_attribute("nuke_repetition", nuke_repetition)
|
||||
|
||||
@set_processing
|
||||
async def converse(self, actor, editor=None):
|
||||
async def converse(self, actor):
|
||||
"""
|
||||
Have a conversation with the AI
|
||||
"""
|
||||
|
||||
history = actor.history
|
||||
self.current_memory_context = None
|
||||
|
||||
character = actor.character
|
||||
|
||||
emission = ConversationAgentEmission(agent=self, generation="", actor=actor, character=character)
|
||||
await talemate.emit.async_signals.get("agent.conversation.before_generate").send(emission)
|
||||
|
||||
|
||||
emission = ConversationAgentEmission(
|
||||
agent=self, generation="", actor=actor, character=character
|
||||
)
|
||||
await talemate.emit.async_signals.get(
|
||||
"agent.conversation.before_generate"
|
||||
).send(emission)
|
||||
|
||||
self.set_generation_overrides()
|
||||
|
||||
result = await self.client.send_prompt(await self.build_prompt(character))
|
||||
@@ -477,7 +600,7 @@ class ConversationAgent(Agent):
|
||||
|
||||
result = self.clean_result(result, character)
|
||||
|
||||
total_result += " "+result
|
||||
total_result += " " + result
|
||||
|
||||
if len(total_result) == 0 and max_loops < 10:
|
||||
max_loops += 1
|
||||
@@ -497,13 +620,19 @@ class ConversationAgent(Agent):
|
||||
|
||||
result = result.replace(" :", ":")
|
||||
|
||||
total_result = total_result.split("#")[0]
|
||||
total_result = total_result.split("#")[0].strip()
|
||||
|
||||
# movie script format
|
||||
# {uppercase character name}
|
||||
# {dialogue}
|
||||
total_result = total_result.replace(f"{character.name.upper()}\n", f"")
|
||||
|
||||
# chat format
|
||||
# {character name}: {dialogue}
|
||||
total_result = total_result.replace(f"{character.name}:", "")
|
||||
|
||||
# Removes partial sentence at the end
|
||||
total_result = util.clean_dialogue(total_result, main_name=character.name)
|
||||
|
||||
# Remove "{character.name}:" - all occurences
|
||||
total_result = total_result.replace(f"{character.name}:", "")
|
||||
|
||||
# Check if total_result starts with character name, if not, prepend it
|
||||
if not total_result.startswith(character.name):
|
||||
@@ -520,13 +649,17 @@ class ConversationAgent(Agent):
|
||||
)
|
||||
|
||||
response_message = util.parse_messages_from_str(total_result, [character.name])
|
||||
|
||||
log.info("conversation agent", result=response_message)
|
||||
|
||||
emission = ConversationAgentEmission(agent=self, generation=response_message, actor=actor, character=character)
|
||||
await talemate.emit.async_signals.get("agent.conversation.generated").send(emission)
|
||||
|
||||
#log.info("conversation agent", generation=emission.generation)
|
||||
log.info("conversation agent", result=response_message)
|
||||
|
||||
emission = ConversationAgentEmission(
|
||||
agent=self, generation=response_message, actor=actor, character=character
|
||||
)
|
||||
await talemate.emit.async_signals.get("agent.conversation.generated").send(
|
||||
emission
|
||||
)
|
||||
|
||||
# log.info("conversation agent", generation=emission.generation)
|
||||
|
||||
messages = [CharacterMessage(message) for message in emission.generation]
|
||||
|
||||
@@ -534,3 +667,18 @@ class ConversationAgent(Agent):
|
||||
actor.scene.push_history(messages)
|
||||
|
||||
return messages
|
||||
|
||||
def allow_repetition_break(
|
||||
self, kind: str, agent_function_name: str, auto: bool = False
|
||||
):
|
||||
if auto and not self.actions["auto_break_repetition"].enabled:
|
||||
return False
|
||||
|
||||
return agent_function_name == "converse"
|
||||
|
||||
def inject_prompt_paramters(
|
||||
self, prompt_param: dict, kind: str, agent_function_name: str
|
||||
):
|
||||
if prompt_param.get("extra_stopping_strings") is None:
|
||||
prompt_param["extra_stopping_strings"] = []
|
||||
prompt_param["extra_stopping_strings"] += ["#"]
|
||||
|
||||
@@ -3,21 +3,23 @@ from __future__ import annotations
|
||||
import json
|
||||
import os
|
||||
|
||||
from talemate.agents.base import Agent
|
||||
import talemate.client as client
|
||||
from talemate.agents.base import Agent, set_processing
|
||||
from talemate.agents.registry import register
|
||||
from talemate.emit import emit
|
||||
import talemate.client as client
|
||||
from talemate.prompts import Prompt
|
||||
|
||||
from .assistant import AssistantMixin
|
||||
from .character import CharacterCreatorMixin
|
||||
from .scenario import ScenarioCreatorMixin
|
||||
|
||||
|
||||
@register()
|
||||
class CreatorAgent(CharacterCreatorMixin, ScenarioCreatorMixin, Agent):
|
||||
|
||||
class CreatorAgent(CharacterCreatorMixin, ScenarioCreatorMixin, AssistantMixin, Agent):
|
||||
"""
|
||||
Creates characters and scenarios and other fun stuff!
|
||||
"""
|
||||
|
||||
|
||||
agent_type = "creator"
|
||||
verbose_name = "Creator"
|
||||
|
||||
@@ -77,12 +79,14 @@ class CreatorAgent(CharacterCreatorMixin, ScenarioCreatorMixin, Agent):
|
||||
# Remove duplicates while preserving the order for list type keys
|
||||
for key, value in merged_data.items():
|
||||
if isinstance(value, list):
|
||||
merged_data[key] = [x for i, x in enumerate(value) if x not in value[:i]]
|
||||
merged_data[key] = [
|
||||
x for i, x in enumerate(value) if x not in value[:i]
|
||||
]
|
||||
|
||||
merged_data["context"] = context
|
||||
|
||||
return merged_data
|
||||
|
||||
|
||||
def load_templates_old(self, names: list, template_type: str = "character") -> dict:
|
||||
"""
|
||||
Loads multiple character creation templates from ./templates/character and merges them in order.
|
||||
@@ -127,8 +131,10 @@ class CreatorAgent(CharacterCreatorMixin, ScenarioCreatorMixin, Agent):
|
||||
|
||||
if "context" in template_data["instructions"]:
|
||||
context = template_data["instructions"]["context"]
|
||||
|
||||
merged_instructions[name]["questions"] = [q[0] for q in template_data.get("questions", [])]
|
||||
|
||||
merged_instructions[name]["questions"] = [
|
||||
q[0] for q in template_data.get("questions", [])
|
||||
]
|
||||
|
||||
# Remove duplicates while preserving the order
|
||||
merged_template = [
|
||||
@@ -157,3 +163,33 @@ class CreatorAgent(CharacterCreatorMixin, ScenarioCreatorMixin, Agent):
|
||||
|
||||
return rv
|
||||
|
||||
@set_processing
|
||||
async def generate_json_list(
|
||||
self,
|
||||
text: str,
|
||||
count: int = 20,
|
||||
first_item: str = None,
|
||||
):
|
||||
_, json_list = await Prompt.request(
|
||||
f"creator.generate-json-list",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"text": text,
|
||||
"first_item": first_item,
|
||||
"count": count,
|
||||
},
|
||||
)
|
||||
return json_list.get("items", [])
|
||||
|
||||
@set_processing
|
||||
async def generate_title(self, text: str):
|
||||
title = await Prompt.request(
|
||||
f"creator.generate-title",
|
||||
self.client,
|
||||
"create_short",
|
||||
vars={
|
||||
"text": text,
|
||||
},
|
||||
)
|
||||
return title
|
||||
|
||||
95
src/talemate/agents/creator/assistant.py
Normal file
@@ -0,0 +1,95 @@
|
||||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
import pydantic
|
||||
|
||||
import talemate.util as util
|
||||
from talemate.agents.base import set_processing
|
||||
from talemate.prompts import Prompt
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate.tale_mate import Character, Scene
|
||||
|
||||
|
||||
class ContentGenerationContext(pydantic.BaseModel):
|
||||
"""
|
||||
A context for generating content.
|
||||
"""
|
||||
|
||||
context: str
|
||||
instructions: str
|
||||
length: int
|
||||
character: Union[str, None] = None
|
||||
original: Union[str, None] = None
|
||||
|
||||
@property
|
||||
def computed_context(self) -> (str, str):
|
||||
typ, context = self.context.split(":", 1)
|
||||
return typ, context
|
||||
|
||||
|
||||
class AssistantMixin:
|
||||
"""
|
||||
Creator mixin that allows quick contextual generation of content.
|
||||
"""
|
||||
|
||||
async def contextual_generate_from_args(
|
||||
self,
|
||||
context: str,
|
||||
instructions: str,
|
||||
length: int = 100,
|
||||
character: Union[str, None] = None,
|
||||
original: Union[str, None] = None,
|
||||
):
|
||||
"""
|
||||
Request content from the assistant.
|
||||
"""
|
||||
|
||||
generation_context = ContentGenerationContext(
|
||||
context=context,
|
||||
instructions=instructions,
|
||||
length=length,
|
||||
character=character,
|
||||
original=original,
|
||||
)
|
||||
|
||||
return await self.contextual_generate(generation_context)
|
||||
|
||||
@set_processing
|
||||
async def contextual_generate(
|
||||
self,
|
||||
generation_context: ContentGenerationContext,
|
||||
):
|
||||
"""
|
||||
Request content from the assistant.
|
||||
"""
|
||||
|
||||
context_typ, context_name = generation_context.computed_context
|
||||
|
||||
if generation_context.length < 100:
|
||||
kind = "create_short"
|
||||
elif generation_context.length < 500:
|
||||
kind = "create_concise"
|
||||
else:
|
||||
kind = "create"
|
||||
|
||||
content = await Prompt.request(
|
||||
f"creator.contextual-generate",
|
||||
self.client,
|
||||
kind,
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"generation_context": generation_context,
|
||||
"context_typ": context_typ,
|
||||
"context_name": context_name,
|
||||
"character": (
|
||||
self.scene.get_character(generation_context.character)
|
||||
if generation_context.character
|
||||
else None
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
content = util.strip_partial_sentences(content)
|
||||
|
||||
return content.strip()
|
||||
@@ -1,42 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import asyncio
|
||||
import random
|
||||
import structlog
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
|
||||
import structlog
|
||||
|
||||
import talemate.util as util
|
||||
from talemate.emit import emit
|
||||
from talemate.prompts import Prompt, LoopedPrompt
|
||||
from talemate.exceptions import LLMAccuracyError
|
||||
from talemate.agents.base import set_processing
|
||||
from talemate.emit import emit
|
||||
from talemate.exceptions import LLMAccuracyError
|
||||
from talemate.prompts import LoopedPrompt, Prompt
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate.tale_mate import Character
|
||||
|
||||
log = structlog.get_logger("talemate.agents.creator.character")
|
||||
|
||||
def validate(k,v):
|
||||
|
||||
def validate(k, v):
|
||||
if k and k.lower() == "gender":
|
||||
return v.lower().strip()
|
||||
if k and k.lower() == "age":
|
||||
try:
|
||||
return int(v.split("\n")[0].strip())
|
||||
except (ValueError, TypeError):
|
||||
raise LLMAccuracyError("Was unable to get a valid age from the response", model_name=None)
|
||||
|
||||
raise LLMAccuracyError(
|
||||
"Was unable to get a valid age from the response", model_name=None
|
||||
)
|
||||
|
||||
return v.strip().strip("\n")
|
||||
|
||||
DEFAULT_CONTENT_CONTEXT="a fun and engaging adventure aimed at an adult audience."
|
||||
|
||||
DEFAULT_CONTENT_CONTEXT = "a fun and engaging adventure aimed at an adult audience."
|
||||
|
||||
|
||||
class CharacterCreatorMixin:
|
||||
"""
|
||||
Adds character creation functionality to the creator agent
|
||||
"""
|
||||
|
||||
|
||||
## NEW
|
||||
|
||||
|
||||
@set_processing
|
||||
async def create_character_attributes(
|
||||
self,
|
||||
@@ -48,8 +54,6 @@ class CharacterCreatorMixin:
|
||||
custom_attributes: dict[str, str] = dict(),
|
||||
predefined_attributes: dict[str, str] = dict(),
|
||||
):
|
||||
|
||||
|
||||
def spice(prompt, spices):
|
||||
# generate number from 0 to 1 and if its smaller than use_spice
|
||||
# select a random spice from the list and return it formatted
|
||||
@@ -57,69 +61,74 @@ class CharacterCreatorMixin:
|
||||
if random.random() < use_spice:
|
||||
spice = random.choice(spices)
|
||||
return prompt.format(spice=spice)
|
||||
return ""
|
||||
|
||||
return ""
|
||||
|
||||
# drop any empty attributes from predefined_attributes
|
||||
|
||||
predefined_attributes = {k:v for k,v in predefined_attributes.items() if v}
|
||||
|
||||
prompt = Prompt.get(f"creator.character-attributes-{template}", vars={
|
||||
"character_prompt": character_prompt,
|
||||
"template": template,
|
||||
"spice": spice,
|
||||
"content_context": content_context,
|
||||
"custom_attributes": custom_attributes,
|
||||
"character_sheet": LoopedPrompt(
|
||||
validate_value=validate,
|
||||
on_update=attribute_callback,
|
||||
generated=predefined_attributes,
|
||||
),
|
||||
})
|
||||
|
||||
predefined_attributes = {k: v for k, v in predefined_attributes.items() if v}
|
||||
|
||||
prompt = Prompt.get(
|
||||
f"creator.character-attributes-{template}",
|
||||
vars={
|
||||
"character_prompt": character_prompt,
|
||||
"template": template,
|
||||
"spice": spice,
|
||||
"content_context": content_context,
|
||||
"custom_attributes": custom_attributes,
|
||||
"character_sheet": LoopedPrompt(
|
||||
validate_value=validate,
|
||||
on_update=attribute_callback,
|
||||
generated=predefined_attributes,
|
||||
),
|
||||
},
|
||||
)
|
||||
await prompt.loop(self.client, "character_sheet", kind="create_concise")
|
||||
|
||||
|
||||
return prompt.vars["character_sheet"].generated
|
||||
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def create_character_description(
|
||||
self,
|
||||
character:Character,
|
||||
self,
|
||||
character: Character,
|
||||
content_context: str = DEFAULT_CONTENT_CONTEXT,
|
||||
):
|
||||
|
||||
description = await Prompt.request(f"creator.character-description", self.client, "create", vars={
|
||||
"character": character,
|
||||
"content_context": content_context,
|
||||
})
|
||||
|
||||
description = await Prompt.request(
|
||||
f"creator.character-description",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"character": character,
|
||||
"content_context": content_context,
|
||||
},
|
||||
)
|
||||
|
||||
return description.strip()
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def create_character_details(
|
||||
self,
|
||||
self,
|
||||
character: Character,
|
||||
template: str,
|
||||
detail_callback: Callable = lambda question, answer: None,
|
||||
questions: list[str] = None,
|
||||
content_context: str = DEFAULT_CONTENT_CONTEXT,
|
||||
):
|
||||
prompt = Prompt.get(f"creator.character-details-{template}", vars={
|
||||
"character_details": LoopedPrompt(
|
||||
validate_value=validate,
|
||||
on_update=detail_callback,
|
||||
),
|
||||
"template": template,
|
||||
"content_context": content_context,
|
||||
"character": character,
|
||||
"custom_questions": questions or [],
|
||||
})
|
||||
prompt = Prompt.get(
|
||||
f"creator.character-details-{template}",
|
||||
vars={
|
||||
"character_details": LoopedPrompt(
|
||||
validate_value=validate,
|
||||
on_update=detail_callback,
|
||||
),
|
||||
"template": template,
|
||||
"content_context": content_context,
|
||||
"character": character,
|
||||
"custom_questions": questions or [],
|
||||
},
|
||||
)
|
||||
await prompt.loop(self.client, "character_details", kind="create_concise")
|
||||
return prompt.vars["character_details"].generated
|
||||
|
||||
|
||||
@set_processing
|
||||
async def create_character_example_dialogue(
|
||||
self,
|
||||
@@ -131,75 +140,135 @@ class CharacterCreatorMixin:
|
||||
example_callback: Callable = lambda example: None,
|
||||
rules_callback: Callable = lambda rules: None,
|
||||
):
|
||||
|
||||
dialogue_rules = await Prompt.request(f"creator.character-dialogue-rules", self.client, "create", vars={
|
||||
"guide": guide,
|
||||
"character": character,
|
||||
"examples": examples or [],
|
||||
"content_context": content_context,
|
||||
})
|
||||
|
||||
dialogue_rules = await Prompt.request(
|
||||
f"creator.character-dialogue-rules",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"guide": guide,
|
||||
"character": character,
|
||||
"examples": examples or [],
|
||||
"content_context": content_context,
|
||||
},
|
||||
)
|
||||
|
||||
log.info("dialogue_rules", dialogue_rules=dialogue_rules)
|
||||
|
||||
|
||||
if rules_callback:
|
||||
rules_callback(dialogue_rules)
|
||||
|
||||
example_dialogue_prompt = Prompt.get(f"creator.character-example-dialogue-{template}", vars={
|
||||
"guide": guide,
|
||||
"character": character,
|
||||
"examples": examples or [],
|
||||
"content_context": content_context,
|
||||
"dialogue_rules": dialogue_rules,
|
||||
"generated_examples": LoopedPrompt(
|
||||
validate_value=validate,
|
||||
on_update=example_callback,
|
||||
),
|
||||
})
|
||||
|
||||
await example_dialogue_prompt.loop(self.client, "generated_examples", kind="create")
|
||||
|
||||
|
||||
example_dialogue_prompt = Prompt.get(
|
||||
f"creator.character-example-dialogue-{template}",
|
||||
vars={
|
||||
"guide": guide,
|
||||
"character": character,
|
||||
"examples": examples or [],
|
||||
"content_context": content_context,
|
||||
"dialogue_rules": dialogue_rules,
|
||||
"generated_examples": LoopedPrompt(
|
||||
validate_value=validate,
|
||||
on_update=example_callback,
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
await example_dialogue_prompt.loop(
|
||||
self.client, "generated_examples", kind="create"
|
||||
)
|
||||
|
||||
return example_dialogue_prompt.vars["generated_examples"].generated
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def determine_content_context_for_character(
|
||||
self,
|
||||
character: Character,
|
||||
):
|
||||
|
||||
content_context = await Prompt.request(f"creator.determine-content-context", self.client, "create", vars={
|
||||
"character": character,
|
||||
})
|
||||
content_context = await Prompt.request(
|
||||
f"creator.determine-content-context",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"character": character,
|
||||
},
|
||||
)
|
||||
return content_context.strip()
|
||||
|
||||
|
||||
@set_processing
|
||||
async def determine_character_attributes(
|
||||
self,
|
||||
character: Character,
|
||||
):
|
||||
|
||||
attributes = await Prompt.request(f"creator.determine-character-attributes", self.client, "analyze_long", vars={
|
||||
"character": character,
|
||||
})
|
||||
attributes = await Prompt.request(
|
||||
f"creator.determine-character-attributes",
|
||||
self.client,
|
||||
"analyze_long",
|
||||
vars={
|
||||
"character": character,
|
||||
},
|
||||
)
|
||||
return attributes
|
||||
|
||||
|
||||
@set_processing
|
||||
async def determine_character_name(
|
||||
self,
|
||||
character_name: str,
|
||||
allowed_names: list[str] = None,
|
||||
) -> str:
|
||||
name = await Prompt.request(
|
||||
f"creator.determine-character-name",
|
||||
self.client,
|
||||
"analyze_freeform_short",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"character_name": character_name,
|
||||
"allowed_names": allowed_names or [],
|
||||
},
|
||||
)
|
||||
return name.split('"', 1)[0].strip().strip(".").strip()
|
||||
|
||||
@set_processing
|
||||
async def determine_character_description(
|
||||
self, character: Character, text: str = ""
|
||||
):
|
||||
description = await Prompt.request(
|
||||
f"creator.determine-character-description",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"character": character,
|
||||
"scene": self.scene,
|
||||
"text": text,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
},
|
||||
)
|
||||
return description.strip()
|
||||
|
||||
@set_processing
|
||||
async def determine_character_goals(
|
||||
self,
|
||||
character: Character,
|
||||
text:str=""
|
||||
goal_instructions: str,
|
||||
):
|
||||
|
||||
description = await Prompt.request(f"creator.determine-character-description", self.client, "create", vars={
|
||||
"character": character,
|
||||
"scene": self.scene,
|
||||
"text": text,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
})
|
||||
return description.strip()
|
||||
|
||||
goals = await Prompt.request(
|
||||
f"creator.determine-character-goals",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"character": character,
|
||||
"scene": self.scene,
|
||||
"goal_instructions": goal_instructions,
|
||||
"npc_name": character.name,
|
||||
"player_name": self.scene.get_player_character().name,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
},
|
||||
)
|
||||
|
||||
log.debug("determine_character_goals", goals=goals, character=character)
|
||||
await character.set_detail("goals", goals.strip())
|
||||
|
||||
return goals.strip()
|
||||
|
||||
@set_processing
|
||||
async def generate_character_from_text(
|
||||
self,
|
||||
@@ -207,11 +276,8 @@ class CharacterCreatorMixin:
|
||||
template: str,
|
||||
content_context: str = DEFAULT_CONTENT_CONTEXT,
|
||||
):
|
||||
|
||||
base_attributes = await self.create_character_attributes(
|
||||
character_prompt=text,
|
||||
template=template,
|
||||
content_context=content_context,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,36 +1,35 @@
|
||||
from talemate.emit import emit, wait_for_input_yesno
|
||||
import re
|
||||
import random
|
||||
import re
|
||||
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.agents.base import set_processing
|
||||
from talemate.emit import emit, wait_for_input_yesno
|
||||
from talemate.prompts import Prompt
|
||||
|
||||
|
||||
class ScenarioCreatorMixin:
|
||||
|
||||
"""
|
||||
Adds scenario creation functionality to the creator agent
|
||||
"""
|
||||
|
||||
|
||||
@set_processing
|
||||
async def create_scene_description(
|
||||
self,
|
||||
prompt:str,
|
||||
content_context:str,
|
||||
prompt: str,
|
||||
content_context: str,
|
||||
):
|
||||
|
||||
"""
|
||||
Creates a new scene.
|
||||
|
||||
|
||||
Arguments:
|
||||
|
||||
|
||||
prompt (str): The prompt to use to create the scene.
|
||||
|
||||
|
||||
content_context (str): The content context to use for the scene.
|
||||
|
||||
|
||||
callback (callable): A callback to call when the scene has been created.
|
||||
"""
|
||||
scene = self.scene
|
||||
|
||||
|
||||
description = await Prompt.request(
|
||||
"creator.scenario-description",
|
||||
self.client,
|
||||
@@ -40,73 +39,70 @@ class ScenarioCreatorMixin:
|
||||
"content_context": content_context,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"scene": scene,
|
||||
}
|
||||
},
|
||||
)
|
||||
description = description.strip()
|
||||
|
||||
|
||||
return description
|
||||
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def create_scene_name(
|
||||
self,
|
||||
prompt:str,
|
||||
content_context:str,
|
||||
description:str,
|
||||
prompt: str,
|
||||
content_context: str,
|
||||
description: str,
|
||||
):
|
||||
|
||||
"""
|
||||
Generates a scene name.
|
||||
|
||||
Arguments:
|
||||
|
||||
prompt (str): The prompt to use to generate the scene name.
|
||||
|
||||
content_context (str): The content context to use for the scene.
|
||||
|
||||
description (str): The description of the scene.
|
||||
"""
|
||||
scene = self.scene
|
||||
|
||||
name = await Prompt.request(
|
||||
"creator.scenario-name",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"prompt": prompt,
|
||||
"content_context": content_context,
|
||||
"description": description,
|
||||
"scene": scene,
|
||||
}
|
||||
)
|
||||
name = name.strip().strip('.!').replace('"','')
|
||||
return name
|
||||
|
||||
|
||||
"""
|
||||
Generates a scene name.
|
||||
|
||||
Arguments:
|
||||
|
||||
prompt (str): The prompt to use to generate the scene name.
|
||||
|
||||
content_context (str): The content context to use for the scene.
|
||||
|
||||
description (str): The description of the scene.
|
||||
"""
|
||||
scene = self.scene
|
||||
|
||||
name = await Prompt.request(
|
||||
"creator.scenario-name",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"prompt": prompt,
|
||||
"content_context": content_context,
|
||||
"description": description,
|
||||
"scene": scene,
|
||||
},
|
||||
)
|
||||
name = name.strip().strip(".!").replace('"', "")
|
||||
return name
|
||||
|
||||
@set_processing
|
||||
async def create_scene_intro(
|
||||
self,
|
||||
prompt:str,
|
||||
content_context:str,
|
||||
description:str,
|
||||
name:str,
|
||||
prompt: str,
|
||||
content_context: str,
|
||||
description: str,
|
||||
name: str,
|
||||
):
|
||||
|
||||
"""
|
||||
Generates a scene introduction.
|
||||
|
||||
|
||||
Arguments:
|
||||
|
||||
|
||||
prompt (str): The prompt to use to generate the scene introduction.
|
||||
|
||||
|
||||
content_context (str): The content context to use for the scene.
|
||||
|
||||
|
||||
description (str): The description of the scene.
|
||||
|
||||
|
||||
name (str): The name of the scene.
|
||||
"""
|
||||
|
||||
|
||||
scene = self.scene
|
||||
|
||||
|
||||
intro = await Prompt.request(
|
||||
"creator.scenario-intro",
|
||||
self.client,
|
||||
@@ -117,17 +113,19 @@ class ScenarioCreatorMixin:
|
||||
"description": description,
|
||||
"name": name,
|
||||
"scene": scene,
|
||||
}
|
||||
},
|
||||
)
|
||||
intro = intro.strip()
|
||||
return intro
|
||||
|
||||
|
||||
@set_processing
|
||||
async def determine_scenario_description(
|
||||
self,
|
||||
text:str
|
||||
):
|
||||
description = await Prompt.request(f"creator.determine-scenario-description", self.client, "analyze_long", vars={
|
||||
"text": text,
|
||||
})
|
||||
async def determine_scenario_description(self, text: str):
|
||||
description = await Prompt.request(
|
||||
f"creator.determine-scenario-description",
|
||||
self.client,
|
||||
"analyze_long",
|
||||
vars={
|
||||
"text": text,
|
||||
},
|
||||
)
|
||||
return description
|
||||
|
||||
34
src/talemate/agents/custom/__init__.py
Normal file
@@ -0,0 +1,34 @@
|
||||
import importlib
|
||||
import os
|
||||
|
||||
import structlog
|
||||
|
||||
log = structlog.get_logger("talemate.agents.custom")
|
||||
|
||||
# import every submodule in this directory
|
||||
#
|
||||
# each directory in this directory is a submodule
|
||||
|
||||
# get the current directory
|
||||
current_directory = os.path.dirname(__file__)
|
||||
|
||||
# get all subdirectories
|
||||
subdirectories = [
|
||||
os.path.join(current_directory, name)
|
||||
for name in os.listdir(current_directory)
|
||||
if os.path.isdir(os.path.join(current_directory, name))
|
||||
]
|
||||
|
||||
# import every submodule
|
||||
|
||||
for subdirectory in subdirectories:
|
||||
# get the name of the submodule
|
||||
submodule_name = os.path.basename(subdirectory)
|
||||
|
||||
if submodule_name.startswith("__"):
|
||||
continue
|
||||
|
||||
log.info("activating custom agent", module=submodule_name)
|
||||
|
||||
# import the submodule
|
||||
importlib.import_module(f".{submodule_name}", __package__)
|
||||
@@ -0,0 +1,5 @@
|
||||
Each agent should be in its own subdirectory.
|
||||
|
||||
The subdirectory itself must be a valid python module.
|
||||
|
||||
Check out docs/dev/agents/example/test for a very simplistic custom agent example.
|
||||
@@ -1,106 +1,384 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import random
|
||||
import structlog
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import talemate.util as util
|
||||
from talemate.emit import wait_for_input, emit
|
||||
import talemate.emit.async_signals
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import NarratorMessage, DirectorMessage
|
||||
from talemate.automated_action import AutomatedAction
|
||||
import structlog
|
||||
|
||||
import talemate.automated_action as automated_action
|
||||
import talemate.emit.async_signals
|
||||
import talemate.instance as instance
|
||||
import talemate.util as util
|
||||
from talemate.agents.conversation import ConversationAgentEmission
|
||||
from talemate.automated_action import AutomatedAction
|
||||
from talemate.emit import emit, wait_for_input
|
||||
from talemate.events import GameLoopActorIterEvent, GameLoopStartEvent, SceneStateEvent
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import DirectorMessage, NarratorMessage
|
||||
|
||||
from .base import Agent, AgentAction, AgentActionConfig, set_processing
|
||||
from .registry import register
|
||||
from .base import set_processing, AgentAction, AgentActionConfig, Agent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate import Actor, Character, Player, Scene
|
||||
|
||||
log = structlog.get_logger("talemate")
|
||||
log = structlog.get_logger("talemate.agent.director")
|
||||
|
||||
|
||||
@register()
|
||||
class DirectorAgent(Agent):
|
||||
agent_type = "director"
|
||||
verbose_name = "Director"
|
||||
|
||||
|
||||
def __init__(self, client, **kwargs):
|
||||
self.is_enabled = False
|
||||
self.is_enabled = True
|
||||
self.client = client
|
||||
self.next_direct = 0
|
||||
self.next_direct_character = {}
|
||||
self.next_direct_scene = 0
|
||||
self.actions = {
|
||||
"direct": AgentAction(enabled=True, label="Direct", description="Will attempt to direct the scene. Runs automatically after AI dialogue (n turns).", config={
|
||||
"turns": AgentActionConfig(type="number", label="Turns", description="Number of turns to wait before directing the sceen", value=5, min=1, max=100, step=1),
|
||||
"prompt": AgentActionConfig(type="text", label="Instructions", description="Instructions to the director", value="", scope="scene")
|
||||
}),
|
||||
"direct": AgentAction(
|
||||
enabled=True,
|
||||
label="Direct",
|
||||
description="Will attempt to direct the scene. Runs automatically after AI dialogue (n turns).",
|
||||
config={
|
||||
"turns": AgentActionConfig(
|
||||
type="number",
|
||||
label="Turns",
|
||||
description="Number of turns to wait before directing the sceen",
|
||||
value=5,
|
||||
min=1,
|
||||
max=100,
|
||||
step=1,
|
||||
),
|
||||
"direct_scene": AgentActionConfig(
|
||||
type="bool",
|
||||
label="Direct Scene",
|
||||
description="If enabled, the scene will be directed through narration",
|
||||
value=True,
|
||||
),
|
||||
"direct_actors": AgentActionConfig(
|
||||
type="bool",
|
||||
label="Direct Actors",
|
||||
description="If enabled, direction will be given to actors based on their goals.",
|
||||
value=True,
|
||||
),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return self.is_enabled
|
||||
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
return True
|
||||
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return True
|
||||
|
||||
|
||||
def connect(self, scene):
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("agent.conversation.before_generate").connect(self.on_conversation_before_generate)
|
||||
|
||||
async def on_conversation_before_generate(self, event:ConversationAgentEmission):
|
||||
talemate.emit.async_signals.get("agent.conversation.before_generate").connect(
|
||||
self.on_conversation_before_generate
|
||||
)
|
||||
talemate.emit.async_signals.get("game_loop_actor_iter").connect(
|
||||
self.on_player_dialog
|
||||
)
|
||||
talemate.emit.async_signals.get("scene_init").connect(self.on_scene_init)
|
||||
|
||||
async def on_scene_init(self, event: SceneStateEvent):
|
||||
"""
|
||||
If game state instructions specify to be run at the start of the game loop
|
||||
we will run them here.
|
||||
"""
|
||||
|
||||
if not self.enabled:
|
||||
if self.scene.game_state.has_scene_instructions:
|
||||
self.is_enabled = True
|
||||
log.warning("on_scene_init - enabling director", scene=self.scene)
|
||||
else:
|
||||
return
|
||||
|
||||
if not self.scene.game_state.has_scene_instructions:
|
||||
return
|
||||
|
||||
if not self.scene.game_state.ops.run_on_start:
|
||||
return
|
||||
|
||||
log.info("on_game_loop_start - running game state instructions")
|
||||
await self.run_gamestate_instructions()
|
||||
|
||||
async def on_conversation_before_generate(self, event: ConversationAgentEmission):
|
||||
log.info("on_conversation_before_generate", director_enabled=self.enabled)
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
await self.direct_scene(event.character)
|
||||
|
||||
async def direct_scene(self, character: Character):
|
||||
|
||||
|
||||
await self.direct(event.character)
|
||||
|
||||
async def on_player_dialog(self, event: GameLoopActorIterEvent):
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
if not self.scene.game_state.has_scene_instructions:
|
||||
return
|
||||
|
||||
if not event.actor.character.is_player:
|
||||
return
|
||||
|
||||
if event.game_loop.had_passive_narration:
|
||||
log.debug(
|
||||
"director.on_player_dialog",
|
||||
skip=True,
|
||||
had_passive_narration=event.game_loop.had_passive_narration,
|
||||
)
|
||||
return
|
||||
|
||||
event.game_loop.had_passive_narration = await self.direct(None)
|
||||
|
||||
async def direct(self, character: Character) -> bool:
|
||||
if not self.actions["direct"].enabled:
|
||||
log.info("direct_scene", skip=True, enabled=self.actions["direct"].enabled)
|
||||
return
|
||||
|
||||
prompt = self.actions["direct"].config["prompt"].value
|
||||
|
||||
if not prompt:
|
||||
log.info("direct_scene", skip=True, prompt=prompt)
|
||||
return
|
||||
|
||||
if self.next_direct % self.actions["direct"].config["turns"].value != 0 or self.next_direct == 0:
|
||||
|
||||
log.info("direct_scene", skip=True, next_direct=self.next_direct)
|
||||
self.next_direct += 1
|
||||
return
|
||||
|
||||
self.next_direct = 0
|
||||
|
||||
await self.direct_character(character, prompt)
|
||||
|
||||
return False
|
||||
|
||||
if character:
|
||||
if not self.actions["direct"].config["direct_actors"].value:
|
||||
log.info(
|
||||
"direct",
|
||||
skip=True,
|
||||
reason="direct_actors disabled",
|
||||
character=character,
|
||||
)
|
||||
return False
|
||||
|
||||
# character direction, see if there are character goals
|
||||
# defined
|
||||
character_goals = character.get_detail("goals")
|
||||
if not character_goals:
|
||||
log.info("direct", skip=True, reason="no goals", character=character)
|
||||
return False
|
||||
|
||||
next_direct = self.next_direct_character.get(character.name, 0)
|
||||
|
||||
if (
|
||||
next_direct % self.actions["direct"].config["turns"].value != 0
|
||||
or next_direct == 0
|
||||
):
|
||||
log.info(
|
||||
"direct", skip=True, next_direct=next_direct, character=character
|
||||
)
|
||||
self.next_direct_character[character.name] = next_direct + 1
|
||||
return False
|
||||
|
||||
self.next_direct_character[character.name] = 0
|
||||
await self.direct_scene(character, character_goals)
|
||||
return True
|
||||
else:
|
||||
if not self.actions["direct"].config["direct_scene"].value:
|
||||
log.info("direct", skip=True, reason="direct_scene disabled")
|
||||
return False
|
||||
|
||||
# no character, see if there are NPC characters at all
|
||||
# if not we always want to direct narration
|
||||
always_direct = (
|
||||
not self.scene.npc_character_names
|
||||
or self.scene.game_state.ops.always_direct
|
||||
)
|
||||
|
||||
next_direct = self.next_direct_scene
|
||||
|
||||
if (
|
||||
next_direct % self.actions["direct"].config["turns"].value != 0
|
||||
or next_direct == 0
|
||||
):
|
||||
if not always_direct:
|
||||
log.info("direct", skip=True, next_direct=next_direct)
|
||||
self.next_direct_scene += 1
|
||||
return False
|
||||
|
||||
self.next_direct_scene = 0
|
||||
await self.direct_scene(None, None)
|
||||
return True
|
||||
|
||||
@set_processing
|
||||
async def direct_character(self, character: Character, prompt:str):
|
||||
|
||||
response = await Prompt.request("director.direct-scene", self.client, "director", vars={
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"scene": self.scene,
|
||||
"prompt": prompt,
|
||||
"character": character,
|
||||
})
|
||||
|
||||
response = response.strip().split("\n")[0].strip()
|
||||
|
||||
response += f" (current story goal: {prompt})"
|
||||
|
||||
log.info("direct_scene", response=response)
|
||||
|
||||
|
||||
message = DirectorMessage(response, source=character.name)
|
||||
emit("director", message, character=character)
|
||||
|
||||
self.scene.push_history(message)
|
||||
async def run_gamestate_instructions(self):
|
||||
"""
|
||||
Run game state instructions, if they exist.
|
||||
"""
|
||||
|
||||
if not self.scene.game_state.has_scene_instructions:
|
||||
return
|
||||
|
||||
await self.direct_scene(None, None)
|
||||
|
||||
@set_processing
|
||||
async def direct_scene(self, character: Character, prompt: str):
|
||||
if not character and self.scene.game_state.game_won:
|
||||
# we are not directing a character, and the game has been won
|
||||
# so we don't need to direct the scene any further
|
||||
return
|
||||
|
||||
if character:
|
||||
# direct a character
|
||||
|
||||
response = await Prompt.request(
|
||||
"director.direct-character",
|
||||
self.client,
|
||||
"director",
|
||||
vars={
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"scene": self.scene,
|
||||
"prompt": prompt,
|
||||
"character": character,
|
||||
"player_character": self.scene.get_player_character(),
|
||||
"game_state": self.scene.game_state,
|
||||
},
|
||||
)
|
||||
|
||||
if "#" in response:
|
||||
response = response.split("#")[0]
|
||||
|
||||
log.info(
|
||||
"direct_character",
|
||||
character=character,
|
||||
prompt=prompt,
|
||||
response=response,
|
||||
)
|
||||
|
||||
response = response.strip().split("\n")[0].strip()
|
||||
# response += f" (current story goal: {prompt})"
|
||||
message = DirectorMessage(response, source=character.name)
|
||||
emit("director", message, character=character)
|
||||
self.scene.push_history(message)
|
||||
else:
|
||||
# run scene instructions
|
||||
self.scene.game_state.scene_instructions
|
||||
|
||||
@set_processing
|
||||
async def persist_characters_from_worldstate(
|
||||
self, exclude: list[str] = None
|
||||
) -> List[Character]:
|
||||
log.warning(
|
||||
"persist_characters_from_worldstate",
|
||||
world_state_characters=self.scene.world_state.characters,
|
||||
scene_characters=self.scene.character_names,
|
||||
)
|
||||
|
||||
created_characters = []
|
||||
|
||||
for character_name in self.scene.world_state.characters.keys():
|
||||
|
||||
if exclude and character_name.lower() in exclude:
|
||||
continue
|
||||
|
||||
if character_name in self.scene.character_names:
|
||||
continue
|
||||
|
||||
character = await self.persist_character(name=character_name)
|
||||
|
||||
created_characters.append(character)
|
||||
|
||||
self.scene.emit_status()
|
||||
|
||||
return created_characters
|
||||
|
||||
@set_processing
|
||||
async def persist_character(
|
||||
self,
|
||||
name: str,
|
||||
content: str = None,
|
||||
attributes: str = None,
|
||||
):
|
||||
world_state = instance.get_agent("world_state")
|
||||
creator = instance.get_agent("creator")
|
||||
|
||||
self.scene.log.debug("persist_character", name=name)
|
||||
name = await creator.determine_character_name(name)
|
||||
self.scene.log.debug("persist_character", adjusted_name=name)
|
||||
|
||||
character = self.scene.Character(name=name)
|
||||
character.color = random.choice(
|
||||
[
|
||||
"#F08080",
|
||||
"#FFD700",
|
||||
"#90EE90",
|
||||
"#ADD8E6",
|
||||
"#DDA0DD",
|
||||
"#FFB6C1",
|
||||
"#FAFAD2",
|
||||
"#D3D3D3",
|
||||
"#B0E0E6",
|
||||
"#FFDEAD",
|
||||
]
|
||||
)
|
||||
|
||||
if not attributes:
|
||||
attributes = await world_state.extract_character_sheet(
|
||||
name=name, text=content
|
||||
)
|
||||
else:
|
||||
attributes = world_state._parse_character_sheet(attributes)
|
||||
|
||||
self.scene.log.debug("persist_character", attributes=attributes)
|
||||
|
||||
character.base_attributes = attributes
|
||||
|
||||
description = await creator.determine_character_description(character)
|
||||
|
||||
character.description = description
|
||||
|
||||
self.scene.log.debug("persist_character", description=description)
|
||||
|
||||
actor = self.scene.Actor(
|
||||
character=character, agent=instance.get_agent("conversation")
|
||||
)
|
||||
|
||||
await self.scene.add_actor(actor)
|
||||
self.scene.emit_status()
|
||||
|
||||
return character
|
||||
|
||||
@set_processing
|
||||
async def update_content_context(
|
||||
self, content: str = None, extra_choices: list[str] = None
|
||||
):
|
||||
if not content:
|
||||
content = "\n".join(
|
||||
self.scene.context_history(sections=False, min_dialogue=25, budget=2048)
|
||||
)
|
||||
|
||||
response = await Prompt.request(
|
||||
"world_state.determine-content-context",
|
||||
self.client,
|
||||
"analyze_freeform",
|
||||
vars={
|
||||
"content": content,
|
||||
"extra_choices": extra_choices or [],
|
||||
},
|
||||
)
|
||||
|
||||
self.scene.context = response.strip()
|
||||
self.scene.emit_status()
|
||||
|
||||
def inject_prompt_paramters(
|
||||
self, prompt_param: dict, kind: str, agent_function_name: str
|
||||
):
|
||||
log.debug(
|
||||
"inject_prompt_paramters",
|
||||
prompt_param=prompt_param,
|
||||
kind=kind,
|
||||
agent_function_name=agent_function_name,
|
||||
)
|
||||
character_names = [f"\n{c.name}:" for c in self.scene.get_characters()]
|
||||
if prompt_param.get("extra_stopping_strings") is None:
|
||||
prompt_param["extra_stopping_strings"] = []
|
||||
prompt_param["extra_stopping_strings"] += character_names + ["#"]
|
||||
if agent_function_name == "update_content_context":
|
||||
prompt_param["extra_stopping_strings"] += ["\n"]
|
||||
|
||||
def allow_repetition_break(
|
||||
self, kind: str, agent_function_name: str, auto: bool = False
|
||||
):
|
||||
return True
|
||||
|
||||
@@ -1,29 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
import traceback
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import structlog
|
||||
|
||||
import talemate.data_objects as data_objects
|
||||
import talemate.util as util
|
||||
import talemate.emit.async_signals
|
||||
import talemate.util as util
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import DirectorMessage, TimePassageMessage
|
||||
|
||||
from .base import Agent, set_processing, AgentAction
|
||||
from .base import Agent, AgentAction, AgentActionConfig, set_processing
|
||||
from .registry import register
|
||||
|
||||
import structlog
|
||||
|
||||
import time
|
||||
import re
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate.tale_mate import Actor, Character, Scene
|
||||
from talemate.agents.conversation import ConversationAgentEmission
|
||||
from talemate.agents.narrator import NarratorAgentEmission
|
||||
from talemate.tale_mate import Actor, Character, Scene
|
||||
|
||||
log = structlog.get_logger("talemate.agents.editor")
|
||||
|
||||
|
||||
@register()
|
||||
class EditorAgent(Agent):
|
||||
"""
|
||||
@@ -34,130 +35,159 @@ class EditorAgent(Agent):
|
||||
|
||||
agent_type = "editor"
|
||||
verbose_name = "Editor"
|
||||
|
||||
|
||||
def __init__(self, client, **kwargs):
|
||||
self.client = client
|
||||
self.is_enabled = True
|
||||
self.actions = {
|
||||
"edit_dialogue": AgentAction(enabled=False, label="Edit dialogue", description="Will attempt to improve the quality of dialogue based on the character and scene. Runs automatically after each AI dialogue."),
|
||||
"fix_exposition": AgentAction(enabled=True, label="Fix exposition", description="Will attempt to fix exposition and emotes, making sure they are displayed in italics. Runs automatically after each AI dialogue."),
|
||||
"add_detail": AgentAction(enabled=False, label="Add detail", description="Will attempt to add extra detail and exposition to the dialogue. Runs automatically after each AI dialogue.")
|
||||
"fix_exposition": AgentAction(
|
||||
enabled=True,
|
||||
label="Fix exposition",
|
||||
description="Will attempt to fix exposition and emotes, making sure they are displayed in italics. Runs automatically after each AI dialogue.",
|
||||
config={
|
||||
"narrator": AgentActionConfig(
|
||||
type="bool",
|
||||
label="Fix narrator messages",
|
||||
description="Will attempt to fix exposition issues in narrator messages",
|
||||
value=True,
|
||||
),
|
||||
},
|
||||
),
|
||||
"add_detail": AgentAction(
|
||||
enabled=False,
|
||||
label="Add detail",
|
||||
description="Will attempt to add extra detail and exposition to the dialogue. Runs automatically after each AI dialogue.",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return self.is_enabled
|
||||
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
return True
|
||||
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return True
|
||||
|
||||
def connect(self, scene):
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("agent.conversation.generated").connect(self.on_conversation_generated)
|
||||
|
||||
async def on_conversation_generated(self, emission:ConversationAgentEmission):
|
||||
talemate.emit.async_signals.get("agent.conversation.generated").connect(
|
||||
self.on_conversation_generated
|
||||
)
|
||||
talemate.emit.async_signals.get("agent.narrator.generated").connect(
|
||||
self.on_narrator_generated
|
||||
)
|
||||
|
||||
async def on_conversation_generated(self, emission: ConversationAgentEmission):
|
||||
"""
|
||||
Called when a conversation is generated
|
||||
"""
|
||||
|
||||
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
|
||||
log.info("editing conversation", emission=emission)
|
||||
|
||||
|
||||
edited = []
|
||||
for text in emission.generation:
|
||||
edit = await self.add_detail(text, emission.character)
|
||||
|
||||
edit = await self.fix_exposition(edit, emission.character)
|
||||
|
||||
|
||||
edit = await self.add_detail(
|
||||
text,
|
||||
emission.character
|
||||
)
|
||||
|
||||
edit = await self.edit_conversation(
|
||||
edit,
|
||||
emission.character
|
||||
)
|
||||
|
||||
edit = await self.fix_exposition(
|
||||
edit,
|
||||
emission.character
|
||||
)
|
||||
|
||||
edited.append(edit)
|
||||
|
||||
|
||||
emission.generation = edited
|
||||
|
||||
|
||||
@set_processing
|
||||
async def edit_conversation(self, content:str, character:Character):
|
||||
|
||||
async def on_narrator_generated(self, emission: NarratorAgentEmission):
|
||||
"""
|
||||
Edits a conversation
|
||||
Called when a narrator message is generated
|
||||
"""
|
||||
|
||||
if not self.actions["edit_dialogue"].enabled:
|
||||
return content
|
||||
|
||||
response = await Prompt.request("editor.edit-dialogue", self.client, "edit_dialogue", vars={
|
||||
"content": content,
|
||||
"character": character,
|
||||
"scene": self.scene,
|
||||
"max_length": self.client.max_token_length
|
||||
})
|
||||
|
||||
response = response.split("[end]")[0]
|
||||
|
||||
response = util.replace_exposition_markers(response)
|
||||
response = util.clean_dialogue(response, main_name=character.name)
|
||||
response = util.strip_partial_sentences(response)
|
||||
|
||||
return response
|
||||
|
||||
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
log.info("editing narrator", emission=emission)
|
||||
|
||||
edited = []
|
||||
|
||||
for text in emission.generation:
|
||||
edit = await self.fix_exposition_on_narrator(text)
|
||||
edited.append(edit)
|
||||
|
||||
emission.generation = edited
|
||||
|
||||
@set_processing
|
||||
async def fix_exposition(self, content:str, character:Character):
|
||||
async def fix_exposition(self, content: str, character: Character):
|
||||
"""
|
||||
Edits a text to make sure all narrative exposition and emotes is encased in *
|
||||
"""
|
||||
|
||||
|
||||
if not self.actions["fix_exposition"].enabled:
|
||||
return content
|
||||
|
||||
#response = await Prompt.request("editor.fix-exposition", self.client, "edit_fix_exposition", vars={
|
||||
# "content": content,
|
||||
# "character": character,
|
||||
# "scene": self.scene,
|
||||
# "max_length": self.client.max_token_length
|
||||
#})
|
||||
|
||||
content = util.clean_dialogue(content, main_name=character.name)
|
||||
|
||||
if not character.is_player:
|
||||
if '"' not in content and "*" not in content:
|
||||
content = util.strip_partial_sentences(content)
|
||||
character_prefix = f"{character.name}: "
|
||||
message = content.split(character_prefix)[1]
|
||||
content = f'{character_prefix}"{message.strip()}"'
|
||||
return content
|
||||
elif '"' in content:
|
||||
# silly hack to clean up some LLMs that always start with a quote
|
||||
# even though the immediate next thing is a narration (indicated by *)
|
||||
content = content.replace(
|
||||
f'{character.name}: "*', f"{character.name}: *"
|
||||
)
|
||||
|
||||
content = util.clean_dialogue(content, main_name=character.name)
|
||||
content = util.strip_partial_sentences(content)
|
||||
content = util.ensure_dialog_format(content, talking_character=character.name)
|
||||
|
||||
|
||||
return content
|
||||
|
||||
|
||||
@set_processing
|
||||
async def add_detail(self, content:str, character:Character):
|
||||
async def fix_exposition_on_narrator(self, content: str):
|
||||
if not self.actions["fix_exposition"].enabled:
|
||||
return content
|
||||
|
||||
if not self.actions["fix_exposition"].config["narrator"].value:
|
||||
return content
|
||||
|
||||
content = util.strip_partial_sentences(content)
|
||||
|
||||
if '"' not in content:
|
||||
content = f"*{content.strip('*')}*"
|
||||
else:
|
||||
content = util.ensure_dialog_format(content)
|
||||
|
||||
return content
|
||||
|
||||
@set_processing
|
||||
async def add_detail(self, content: str, character: Character):
|
||||
"""
|
||||
Edits a text to increase its length and add extra detail and exposition
|
||||
"""
|
||||
|
||||
|
||||
if not self.actions["add_detail"].enabled:
|
||||
return content
|
||||
|
||||
response = await Prompt.request("editor.add-detail", self.client, "edit_add_detail", vars={
|
||||
"content": content,
|
||||
"character": character,
|
||||
"scene": self.scene,
|
||||
"max_length": self.client.max_token_length
|
||||
})
|
||||
|
||||
|
||||
response = await Prompt.request(
|
||||
"editor.add-detail",
|
||||
self.client,
|
||||
"edit_add_detail",
|
||||
vars={
|
||||
"content": content,
|
||||
"character": character,
|
||||
"scene": self.scene,
|
||||
"max_length": self.client.max_token_length,
|
||||
},
|
||||
)
|
||||
|
||||
response = util.replace_exposition_markers(response)
|
||||
response = util.clean_dialogue(response, main_name=character.name)
|
||||
response = util.clean_dialogue(response, main_name=character.name)
|
||||
response = util.strip_partial_sentences(response)
|
||||
|
||||
return response
|
||||
|
||||
return response
|
||||
|
||||
@@ -1,15 +1,21 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import os
|
||||
import shutil
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import structlog
|
||||
from chromadb.config import Settings
|
||||
|
||||
import talemate.events as events
|
||||
import talemate.util as util
|
||||
from talemate.context import scene_is_loading
|
||||
from talemate.agents.base import set_processing
|
||||
from talemate.config import load_config
|
||||
import structlog
|
||||
import shutil
|
||||
from talemate.context import scene_is_loading
|
||||
from talemate.emit import emit
|
||||
from talemate.emit.signals import handlers
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
@@ -24,7 +30,18 @@ if not chromadb:
|
||||
log.info("ChromaDB not found, disabling Chroma agent")
|
||||
|
||||
|
||||
from .base import Agent
|
||||
from .base import Agent, AgentDetail
|
||||
|
||||
|
||||
class MemoryDocument(str):
|
||||
def __new__(cls, text, meta, id, raw):
|
||||
inst = super().__new__(cls, text)
|
||||
|
||||
inst.meta = meta
|
||||
inst.id = id
|
||||
inst.raw = raw
|
||||
|
||||
return inst
|
||||
|
||||
|
||||
class MemoryAgent(Agent):
|
||||
@@ -38,10 +55,11 @@ class MemoryAgent(Agent):
|
||||
|
||||
@property
|
||||
def readonly(self):
|
||||
|
||||
if scene_is_loading.get() and not getattr(self.scene, "_memory_never_persisted", False):
|
||||
if scene_is_loading.get() and not getattr(
|
||||
self.scene, "_memory_never_persisted", False
|
||||
):
|
||||
return True
|
||||
|
||||
|
||||
return False
|
||||
|
||||
@property
|
||||
@@ -57,6 +75,16 @@ class MemoryAgent(Agent):
|
||||
self.scene = scene
|
||||
self.memory_tracker = {}
|
||||
self.config = load_config()
|
||||
self._ready_to_add = False
|
||||
|
||||
handlers["config_saved"].connect(self.on_config_saved)
|
||||
|
||||
def on_config_saved(self, event):
|
||||
openai_key = self.openai_api_key
|
||||
self.config = load_config()
|
||||
if openai_key != self.openai_api_key:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.run_until_complete(self.emit_status())
|
||||
|
||||
async def set_db(self):
|
||||
raise NotImplementedError()
|
||||
@@ -67,40 +95,135 @@ class MemoryAgent(Agent):
|
||||
async def count(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
async def add(self, text, character=None, uid=None, ts:str=None, **kwargs):
|
||||
@set_processing
|
||||
async def add(self, text, character=None, uid=None, ts: str = None, **kwargs):
|
||||
if not text:
|
||||
return
|
||||
if self.readonly:
|
||||
log.debug("memory agent", status="readonly")
|
||||
return
|
||||
await self._add(text, character=character, uid=uid, ts=ts, **kwargs)
|
||||
|
||||
async def _add(self, text, character=None, ts:str=None, **kwargs):
|
||||
while not self._ready_to_add:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
log.debug(
|
||||
"memory agent add",
|
||||
text=text[:50],
|
||||
character=character,
|
||||
uid=uid,
|
||||
ts=ts,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
functools.partial(self._add, text, character, uid=uid, ts=ts, **kwargs),
|
||||
)
|
||||
except AttributeError as e:
|
||||
# not sure how this sometimes happens.
|
||||
# chromadb model None
|
||||
# race condition because we are forcing async context onto it?
|
||||
|
||||
log.error(
|
||||
"memory agent",
|
||||
error="failed to add memory",
|
||||
details=e,
|
||||
text=text[:50],
|
||||
character=character,
|
||||
uid=uid,
|
||||
ts=ts,
|
||||
**kwargs,
|
||||
)
|
||||
await asyncio.sleep(1.0)
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
functools.partial(
|
||||
self._add, text, character, uid=uid, ts=ts, **kwargs
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
log.error(
|
||||
"memory agent",
|
||||
error="failed to add memory (retried)",
|
||||
details=e,
|
||||
text=text[:50],
|
||||
character=character,
|
||||
uid=uid,
|
||||
ts=ts,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def _add(self, text, character=None, ts: str = None, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
@set_processing
|
||||
async def add_many(self, objects: list[dict]):
|
||||
if self.readonly:
|
||||
log.debug("memory agent", status="readonly")
|
||||
return
|
||||
await self._add_many(objects)
|
||||
|
||||
async def _add_many(self, objects: list[dict]):
|
||||
|
||||
while not self._ready_to_add:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
log.debug("memory agent add many", len=len(objects))
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self._add_many, objects)
|
||||
|
||||
def _add_many(self, objects: list[dict]):
|
||||
"""
|
||||
Add multiple objects to the memory
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def get(self, text, character=None, **query):
|
||||
return await self._get(str(text), character, **query)
|
||||
|
||||
async def _get(self, text, character=None, **query):
|
||||
def _delete(self, meta: dict):
|
||||
"""
|
||||
Delete an object from the memory
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_document(self, id):
|
||||
return self.db.get(id)
|
||||
@set_processing
|
||||
async def delete(self, meta: dict):
|
||||
"""
|
||||
Delete an object from the memory
|
||||
"""
|
||||
if self.readonly:
|
||||
log.debug("memory agent", status="readonly")
|
||||
return
|
||||
|
||||
while not self._ready_to_add:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self._delete, meta)
|
||||
|
||||
@set_processing
|
||||
async def get(self, text, character=None, **query):
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
return await loop.run_in_executor(
|
||||
None, functools.partial(self._get, text, character, **query)
|
||||
)
|
||||
|
||||
def _get(self, text, character=None, **query):
|
||||
raise NotImplementedError()
|
||||
|
||||
@set_processing
|
||||
async def get_document(self, id):
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(None, self._get_document, id)
|
||||
|
||||
def _get_document(self, id):
|
||||
raise NotImplementedError()
|
||||
|
||||
def on_archive_add(self, event: events.ArchiveEvent):
|
||||
asyncio.ensure_future(self.add(event.text, uid=event.memory_id, ts=event.ts, typ="history"))
|
||||
asyncio.ensure_future(
|
||||
self.add(event.text, uid=event.memory_id, ts=event.ts, typ="history")
|
||||
)
|
||||
|
||||
def on_character_state(self, event: events.CharacterStateEvent):
|
||||
asyncio.ensure_future(
|
||||
@@ -140,6 +263,10 @@ class MemoryAgent(Agent):
|
||||
"""
|
||||
|
||||
memory_context = []
|
||||
|
||||
if not query:
|
||||
return memory_context
|
||||
|
||||
for memory in await self.get(query):
|
||||
if memory in memory_context:
|
||||
continue
|
||||
@@ -153,17 +280,26 @@ class MemoryAgent(Agent):
|
||||
break
|
||||
return memory_context
|
||||
|
||||
async def query(self, query:str, max_tokens:int=1000, filter:Callable=lambda x:True, **where):
|
||||
async def query(
|
||||
self,
|
||||
query: str,
|
||||
max_tokens: int = 1000,
|
||||
filter: Callable = lambda x: True,
|
||||
**where,
|
||||
):
|
||||
"""
|
||||
Get the character memory context for a given character
|
||||
"""
|
||||
|
||||
try:
|
||||
return (await self.multi_query([query], max_tokens=max_tokens, filter=filter, **where))[0]
|
||||
return (
|
||||
await self.multi_query(
|
||||
[query], max_tokens=max_tokens, filter=filter, **where
|
||||
)
|
||||
)[0]
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
|
||||
async def multi_query(
|
||||
self,
|
||||
queries: list[str],
|
||||
@@ -171,7 +307,8 @@ class MemoryAgent(Agent):
|
||||
max_tokens: int = 1000,
|
||||
filter: Callable = lambda x: True,
|
||||
formatter: Callable = lambda x: x,
|
||||
**where
|
||||
limit: int = 10,
|
||||
**where,
|
||||
):
|
||||
"""
|
||||
Get the character memory context for a given character
|
||||
@@ -179,8 +316,11 @@ class MemoryAgent(Agent):
|
||||
|
||||
memory_context = []
|
||||
for query in queries:
|
||||
if not query:
|
||||
continue
|
||||
|
||||
i = 0
|
||||
for memory in await self.get(formatter(query), limit=iterate, **where):
|
||||
for memory in await self.get(formatter(query), limit=limit, **where):
|
||||
if memory in memory_context:
|
||||
continue
|
||||
|
||||
@@ -205,10 +345,13 @@ from .registry import register
|
||||
|
||||
@register(condition=lambda: chromadb is not None)
|
||||
class ChromaDBMemoryAgent(MemoryAgent):
|
||||
|
||||
requires_llm_client = False
|
||||
|
||||
@property
|
||||
def ready(self):
|
||||
if self.embeddings == "openai" and not self.openai_api_key:
|
||||
return False
|
||||
|
||||
if getattr(self, "db_client", None):
|
||||
return True
|
||||
return False
|
||||
@@ -217,77 +360,125 @@ class ChromaDBMemoryAgent(MemoryAgent):
|
||||
def status(self):
|
||||
if self.ready:
|
||||
return "active" if not getattr(self, "processing", False) else "busy"
|
||||
|
||||
if self.embeddings == "openai" and not self.openai_api_key:
|
||||
return "error"
|
||||
|
||||
return "waiting"
|
||||
|
||||
@property
|
||||
def agent_details(self):
|
||||
|
||||
details = {
|
||||
"backend": AgentDetail(
|
||||
icon="mdi-server-outline",
|
||||
value="ChromaDB",
|
||||
description="The backend to use for long-term memory",
|
||||
).model_dump(),
|
||||
"embeddings": AgentDetail(
|
||||
icon="mdi-cube-unfolded",
|
||||
value=self.embeddings,
|
||||
description="The embeddings model.",
|
||||
).model_dump(),
|
||||
}
|
||||
|
||||
if self.embeddings == "openai" and not self.openai_api_key:
|
||||
# return "No OpenAI API key set"
|
||||
details["error"] = {
|
||||
"icon": "mdi-alert",
|
||||
"value": "No OpenAI API key set",
|
||||
"description": "You must provide an OpenAI API key to use OpenAI embeddings",
|
||||
"color": "error",
|
||||
}
|
||||
|
||||
return details
|
||||
|
||||
return f"ChromaDB: {self.embeddings}"
|
||||
|
||||
@property
|
||||
def embeddings(self):
|
||||
"""
|
||||
Returns which embeddings to use
|
||||
|
||||
|
||||
will read from TM_CHROMADB_EMBEDDINGS env variable and default to 'default' using
|
||||
the default embeddings specified by chromadb.
|
||||
|
||||
|
||||
other values are
|
||||
|
||||
|
||||
- openai: use openai embeddings
|
||||
- instructor: use instructor embeddings
|
||||
|
||||
|
||||
for `openai`:
|
||||
|
||||
|
||||
you will also need to provide an `OPENAI_API_KEY` env variable
|
||||
|
||||
|
||||
for `instructor`:
|
||||
|
||||
|
||||
you will also need to provide which instructor model to use with the `TM_INSTRUCTOR_MODEL` env variable, which defaults to hkunlp/instructor-xl
|
||||
|
||||
|
||||
additionally you can provide the `TM_INSTRUCTOR_DEVICE` env variable to specify which device to use, which defaults to cpu
|
||||
"""
|
||||
|
||||
|
||||
embeddings = self.config.get("chromadb").get("embeddings")
|
||||
|
||||
assert embeddings in ["default", "openai", "instructor"], f"Unknown embeddings {embeddings}"
|
||||
|
||||
|
||||
assert embeddings in [
|
||||
"default",
|
||||
"openai",
|
||||
"instructor",
|
||||
], f"Unknown embeddings {embeddings}"
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
@property
|
||||
def USE_OPENAI(self):
|
||||
return self.embeddings == "openai"
|
||||
|
||||
|
||||
@property
|
||||
def USE_INSTRUCTOR(self):
|
||||
return self.embeddings == "instructor"
|
||||
|
||||
|
||||
@property
|
||||
def db_name(self):
|
||||
return getattr(self, "collection_name", "<unnamed>")
|
||||
|
||||
@property
|
||||
def openai_api_key(self):
|
||||
return self.config.get("openai", {}).get("api_key")
|
||||
|
||||
def make_collection_name(self, scene):
|
||||
|
||||
if self.USE_OPENAI:
|
||||
suffix = "-openai"
|
||||
model_name = self.config.get("chromadb").get(
|
||||
"openai_model", "text-embedding-3-small"
|
||||
)
|
||||
if model_name == "text-embedding-ada-002":
|
||||
suffix = "-openai"
|
||||
else:
|
||||
suffix = f"-openai-{model_name}"
|
||||
elif self.USE_INSTRUCTOR:
|
||||
suffix = "-instructor"
|
||||
model = self.config.get("chromadb").get("instructor_model", "hkunlp/instructor-xl")
|
||||
model = self.config.get("chromadb").get(
|
||||
"instructor_model", "hkunlp/instructor-xl"
|
||||
)
|
||||
if "xl" in model:
|
||||
suffix += "-xl"
|
||||
elif "large" in model:
|
||||
suffix += "-large"
|
||||
else:
|
||||
suffix = ""
|
||||
|
||||
|
||||
return f"{scene.memory_id}-tm{suffix}"
|
||||
|
||||
async def count(self):
|
||||
await asyncio.sleep(0)
|
||||
return self.db.count()
|
||||
|
||||
@set_processing
|
||||
async def set_db(self):
|
||||
await self.emit_status(processing=True)
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self._set_db)
|
||||
|
||||
def _set_db(self):
|
||||
self._ready_to_add = False
|
||||
|
||||
if not getattr(self, "db_client", None):
|
||||
log.info("chromadb agent", status="setting up db client to persistent db")
|
||||
@@ -295,69 +486,90 @@ class ChromaDBMemoryAgent(MemoryAgent):
|
||||
settings=Settings(anonymized_telemetry=False)
|
||||
)
|
||||
|
||||
openai_key = self.config.get("openai").get("api_key") or os.environ.get("OPENAI_API_KEY")
|
||||
|
||||
openai_key = self.openai_api_key
|
||||
|
||||
self.collection_name = collection_name = self.make_collection_name(self.scene)
|
||||
|
||||
log.info("chromadb agent", status="setting up db", collection_name=collection_name)
|
||||
|
||||
|
||||
log.info(
|
||||
"chromadb agent", status="setting up db", collection_name=collection_name
|
||||
)
|
||||
|
||||
if self.USE_OPENAI:
|
||||
|
||||
if not openai_key:
|
||||
raise ValueError("You must provide an the openai ai key in the config if you want to use it for chromadb embeddings")
|
||||
|
||||
raise ValueError(
|
||||
"You must provide an the openai ai key in the config if you want to use it for chromadb embeddings"
|
||||
)
|
||||
|
||||
model_name = self.config.get("chromadb").get(
|
||||
"openai_model", "text-embedding-3-small"
|
||||
)
|
||||
|
||||
log.info(
|
||||
"crhomadb", status="using openai", openai_key=openai_key[:5] + "..."
|
||||
"crhomadb",
|
||||
status="using openai",
|
||||
openai_key=openai_key[:5] + "...",
|
||||
model=model_name,
|
||||
)
|
||||
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
|
||||
api_key = openai_key,
|
||||
model_name="text-embedding-ada-002",
|
||||
api_key=openai_key,
|
||||
model_name=model_name,
|
||||
)
|
||||
self.db = self.db_client.get_or_create_collection(
|
||||
collection_name, embedding_function=openai_ef
|
||||
)
|
||||
elif self.USE_INSTRUCTOR:
|
||||
|
||||
instructor_device = self.config.get("chromadb").get("instructor_device", "cpu")
|
||||
instructor_model = self.config.get("chromadb").get("instructor_model", "hkunlp/instructor-xl")
|
||||
|
||||
log.info("chromadb", status="using instructor", model=instructor_model, device=instructor_device)
|
||||
|
||||
instructor_device = self.config.get("chromadb").get(
|
||||
"instructor_device", "cpu"
|
||||
)
|
||||
instructor_model = self.config.get("chromadb").get(
|
||||
"instructor_model", "hkunlp/instructor-xl"
|
||||
)
|
||||
|
||||
log.info(
|
||||
"chromadb",
|
||||
status="using instructor",
|
||||
model=instructor_model,
|
||||
device=instructor_device,
|
||||
)
|
||||
|
||||
# ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-mpnet-base-v2")
|
||||
ef = embedding_functions.InstructorEmbeddingFunction(
|
||||
model_name=instructor_model, device=instructor_device
|
||||
)
|
||||
|
||||
|
||||
log.info("chromadb", status="embedding function ready")
|
||||
|
||||
|
||||
self.db = self.db_client.get_or_create_collection(
|
||||
collection_name, embedding_function=ef
|
||||
)
|
||||
|
||||
|
||||
log.info("chromadb", status="instructor db ready")
|
||||
else:
|
||||
log.info("chromadb", status="using default embeddings")
|
||||
self.db = self.db_client.get_or_create_collection(collection_name)
|
||||
|
||||
self.scene._memory_never_persisted = self.db.count() == 0
|
||||
|
||||
await self.emit_status(processing=False)
|
||||
self.scene._memory_never_persisted = self.db.count() == 0
|
||||
log.info("chromadb agent", status="db ready")
|
||||
self._ready_to_add = True
|
||||
|
||||
def clear_db(self):
|
||||
if not self.db:
|
||||
return
|
||||
|
||||
log.info("chromadb agent", status="clearing db", collection_name=self.collection_name)
|
||||
|
||||
|
||||
log.info(
|
||||
"chromadb agent", status="clearing db", collection_name=self.collection_name
|
||||
)
|
||||
|
||||
self.db.delete(where={"source": "talemate"})
|
||||
|
||||
|
||||
def drop_db(self):
|
||||
if not self.db:
|
||||
return
|
||||
|
||||
log.info("chromadb agent", status="dropping db", collection_name=self.collection_name)
|
||||
|
||||
|
||||
log.info(
|
||||
"chromadb agent", status="dropping db", collection_name=self.collection_name
|
||||
)
|
||||
|
||||
try:
|
||||
self.db_client.delete_collection(self.collection_name)
|
||||
except ValueError as exc:
|
||||
@@ -367,29 +579,43 @@ class ChromaDBMemoryAgent(MemoryAgent):
|
||||
def close_db(self, scene):
|
||||
if not self.db:
|
||||
return
|
||||
|
||||
log.info("chromadb agent", status="closing db", collection_name=self.collection_name)
|
||||
|
||||
if not scene.saved:
|
||||
|
||||
log.info(
|
||||
"chromadb agent", status="closing db", collection_name=self.collection_name
|
||||
)
|
||||
|
||||
if not scene.saved and not scene.saved_memory_session_id:
|
||||
# scene was never saved so we can discard the memory
|
||||
collection_name = self.make_collection_name(scene)
|
||||
log.info("chromadb agent", status="discarding memory", collection_name=collection_name)
|
||||
log.info(
|
||||
"chromadb agent",
|
||||
status="discarding memory",
|
||||
collection_name=collection_name,
|
||||
)
|
||||
try:
|
||||
self.db_client.delete_collection(collection_name)
|
||||
except ValueError as exc:
|
||||
if "Collection not found" not in str(exc):
|
||||
raise
|
||||
|
||||
log.error(
|
||||
"chromadb agent", error="failed to delete collection", details=exc
|
||||
)
|
||||
elif not scene.saved:
|
||||
# scene was saved but memory was never persisted
|
||||
# so we need to remove the memory from the db
|
||||
self._remove_unsaved_memory()
|
||||
|
||||
self.db = None
|
||||
|
||||
async def _add(self, text, character=None, uid=None, ts:str=None, **kwargs):
|
||||
|
||||
def _add(self, text, character=None, uid=None, ts: str = None, **kwargs):
|
||||
metadatas = []
|
||||
ids = []
|
||||
|
||||
await self.emit_status(processing=True)
|
||||
scene = self.scene
|
||||
|
||||
if character:
|
||||
meta = {"character": character.name, "source": "talemate"}
|
||||
meta = {
|
||||
"character": character.name,
|
||||
"source": "talemate",
|
||||
"session": scene.memory_session_id,
|
||||
}
|
||||
if ts:
|
||||
meta["ts"] = ts
|
||||
meta.update(kwargs)
|
||||
@@ -399,7 +625,11 @@ class ChromaDBMemoryAgent(MemoryAgent):
|
||||
id = uid or f"{character.name}-{self.memory_tracker[character.name]}"
|
||||
ids = [id]
|
||||
else:
|
||||
meta = {"character": "__narrator__", "source": "talemate"}
|
||||
meta = {
|
||||
"character": "__narrator__",
|
||||
"source": "talemate",
|
||||
"session": scene.memory_session_id,
|
||||
}
|
||||
if ts:
|
||||
meta["ts"] = ts
|
||||
meta.update(kwargs)
|
||||
@@ -409,83 +639,104 @@ class ChromaDBMemoryAgent(MemoryAgent):
|
||||
id = uid or f"__narrator__-{self.memory_tracker['__narrator__']}"
|
||||
ids = [id]
|
||||
|
||||
log.debug("chromadb agent add", text=text, meta=meta, id=id)
|
||||
# log.debug("chromadb agent add", text=text, meta=meta, id=id)
|
||||
|
||||
self.db.upsert(documents=[text], metadatas=metadatas, ids=ids)
|
||||
|
||||
await self.emit_status(processing=False)
|
||||
|
||||
async def _add_many(self, objects: list[dict]):
|
||||
|
||||
def _add_many(self, objects: list[dict]):
|
||||
documents = []
|
||||
metadatas = []
|
||||
ids = []
|
||||
scene = self.scene
|
||||
|
||||
await self.emit_status(processing=True)
|
||||
if not objects:
|
||||
return
|
||||
|
||||
for obj in objects:
|
||||
documents.append(obj["text"])
|
||||
meta = obj.get("meta", {})
|
||||
source = meta.get("source", "talemate")
|
||||
character = meta.get("character", "__narrator__")
|
||||
self.memory_tracker.setdefault(character, 0)
|
||||
self.memory_tracker[character] += 1
|
||||
meta["source"] = "talemate"
|
||||
meta["source"] = source
|
||||
if not meta.get("session"):
|
||||
meta["session"] = scene.memory_session_id
|
||||
metadatas.append(meta)
|
||||
uid = obj.get("id", f"{character}-{self.memory_tracker[character]}")
|
||||
ids.append(uid)
|
||||
self.db.upsert(documents=documents, metadatas=metadatas, ids=ids)
|
||||
|
||||
await self.emit_status(processing=False)
|
||||
def _delete(self, meta: dict):
|
||||
if "ids" in meta:
|
||||
log.debug("chromadb agent delete", ids=meta["ids"])
|
||||
self.db.delete(ids=meta["ids"])
|
||||
return
|
||||
|
||||
async def _get(self, text, character=None, limit:int=15, **kwargs):
|
||||
await self.emit_status(processing=True)
|
||||
where = {"$and": [{k: v} for k, v in meta.items()]}
|
||||
self.db.delete(where=where)
|
||||
log.debug("chromadb agent delete", meta=meta, where=where)
|
||||
|
||||
def _get(self, text, character=None, limit: int = 15, **kwargs):
|
||||
where = {}
|
||||
|
||||
# this doesn't work because chromadb currently doesn't match
|
||||
# non existing fields with $ne (or so it seems)
|
||||
# where.setdefault("$and", [{"pin_only": {"$ne": True}}])
|
||||
|
||||
where.setdefault("$and", [])
|
||||
|
||||
|
||||
character_filtered = False
|
||||
|
||||
for k,v in kwargs.items():
|
||||
|
||||
for k, v in kwargs.items():
|
||||
if k == "character":
|
||||
character_filtered = True
|
||||
where["$and"].append({k: v})
|
||||
|
||||
|
||||
if character and not character_filtered:
|
||||
where["$and"].append({"character": character.name})
|
||||
|
||||
|
||||
if len(where["$and"]) == 1:
|
||||
where = where["$and"][0]
|
||||
elif not where["$and"]:
|
||||
where = None
|
||||
|
||||
#log.debug("crhomadb agent get", text=text, where=where)
|
||||
# log.debug("crhomadb agent get", text=text, where=where)
|
||||
|
||||
_results = self.db.query(query_texts=[text], where=where, n_results=limit)
|
||||
|
||||
#import json
|
||||
#print(json.dumps(_results["ids"], indent=2))
|
||||
#print(json.dumps(_results["distances"], indent=2))
|
||||
|
||||
|
||||
# import json
|
||||
# print(json.dumps(_results["ids"], indent=2))
|
||||
# print(json.dumps(_results["distances"], indent=2))
|
||||
|
||||
results = []
|
||||
|
||||
max_distance = 1.5
|
||||
if self.USE_INSTRUCTOR:
|
||||
max_distance = 1
|
||||
elif self.USE_OPENAI:
|
||||
max_distance = 1
|
||||
|
||||
for i in range(len(_results["distances"][0])):
|
||||
distance = _results["distances"][0][i]
|
||||
|
||||
|
||||
doc = _results["documents"][0][i]
|
||||
meta = _results["metadatas"][0][i]
|
||||
ts = meta.get("ts")
|
||||
|
||||
if distance < 1:
|
||||
|
||||
try:
|
||||
date_prefix = util.iso8601_diff_to_human(ts, self.scene.ts)
|
||||
except Exception:
|
||||
log.error("chromadb agent", error="failed to get date prefix", ts=ts, scene_ts=self.scene.ts)
|
||||
date_prefix = None
|
||||
|
||||
|
||||
# skip pin_only entries
|
||||
if meta.get("pin_only", False):
|
||||
continue
|
||||
|
||||
if distance < max_distance:
|
||||
date_prefix = self.convert_ts_to_date_prefix(ts)
|
||||
raw = doc
|
||||
|
||||
if date_prefix:
|
||||
doc = f"{date_prefix}: {doc}"
|
||||
|
||||
doc = MemoryDocument(doc, meta, _results["ids"][0][i], raw)
|
||||
|
||||
results.append(doc)
|
||||
else:
|
||||
break
|
||||
@@ -495,6 +746,57 @@ class ChromaDBMemoryAgent(MemoryAgent):
|
||||
if len(results) > limit:
|
||||
break
|
||||
|
||||
await self.emit_status(processing=False)
|
||||
|
||||
return results
|
||||
|
||||
def convert_ts_to_date_prefix(self, ts):
|
||||
if not ts:
|
||||
return None
|
||||
try:
|
||||
return util.iso8601_diff_to_human(ts, self.scene.ts)
|
||||
except Exception as e:
|
||||
log.error(
|
||||
"chromadb agent",
|
||||
error="failed to get date prefix",
|
||||
details=e,
|
||||
ts=ts,
|
||||
scene_ts=self.scene.ts,
|
||||
)
|
||||
return None
|
||||
|
||||
def _get_document(self, id) -> dict:
|
||||
result = self.db.get(ids=[id] if isinstance(id, str) else id)
|
||||
documents = {}
|
||||
|
||||
for idx, doc in enumerate(result["documents"]):
|
||||
date_prefix = self.convert_ts_to_date_prefix(
|
||||
result["metadatas"][idx].get("ts")
|
||||
)
|
||||
if date_prefix:
|
||||
doc = f"{date_prefix}: {doc}"
|
||||
documents[result["ids"][idx]] = MemoryDocument(
|
||||
doc, result["metadatas"][idx], result["ids"][idx], doc
|
||||
)
|
||||
|
||||
return documents
|
||||
|
||||
@set_processing
|
||||
async def remove_unsaved_memory(self):
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self._remove_unsaved_memory)
|
||||
|
||||
def _remove_unsaved_memory(self):
|
||||
scene = self.scene
|
||||
|
||||
if not scene.memory_session_id:
|
||||
return
|
||||
|
||||
if scene.saved_memory_session_id == self.scene.memory_session_id:
|
||||
return
|
||||
|
||||
log.info(
|
||||
"chromadb agent",
|
||||
status="removing unsaved memory",
|
||||
session_id=scene.memory_session_id,
|
||||
)
|
||||
|
||||
self._delete({"session": scene.memory_session_id, "source": "talemate"})
|
||||
|
||||
@@ -1,87 +1,160 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
import structlog
|
||||
import dataclasses
|
||||
import random
|
||||
import talemate.util as util
|
||||
from talemate.emit import emit
|
||||
import talemate.emit.async_signals
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.agents.base import set_processing, Agent, AgentAction, AgentActionConfig
|
||||
from talemate.agents.world_state import TimePassageEmission
|
||||
from talemate.scene_message import NarratorMessage
|
||||
from talemate.events import GameLoopActorIterEvent
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import structlog
|
||||
|
||||
import talemate.client as client
|
||||
import talemate.emit.async_signals
|
||||
import talemate.util as util
|
||||
from talemate.agents.base import Agent, AgentAction, AgentActionConfig, AgentEmission
|
||||
from talemate.agents.base import set_processing as _set_processing
|
||||
from talemate.agents.world_state import TimePassageEmission
|
||||
from talemate.emit import emit
|
||||
from talemate.events import GameLoopActorIterEvent
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import NarratorMessage
|
||||
|
||||
from .registry import register
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate.tale_mate import Actor, Player, Character
|
||||
from talemate.tale_mate import Actor, Character, Player
|
||||
|
||||
log = structlog.get_logger("talemate.agents.narrator")
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class NarratorAgentEmission(AgentEmission):
|
||||
generation: list[str] = dataclasses.field(default_factory=list)
|
||||
|
||||
|
||||
talemate.emit.async_signals.register("agent.narrator.generated")
|
||||
|
||||
|
||||
def set_processing(fn):
|
||||
"""
|
||||
Custom decorator that emits the agent status as processing while the function
|
||||
is running and then emits the result of the function as a NarratorAgentEmission
|
||||
"""
|
||||
|
||||
@_set_processing
|
||||
@wraps(fn)
|
||||
async def narration_wrapper(self, *args, **kwargs):
|
||||
response = await fn(self, *args, **kwargs)
|
||||
emission = NarratorAgentEmission(
|
||||
agent=self,
|
||||
generation=[response],
|
||||
)
|
||||
await talemate.emit.async_signals.get("agent.narrator.generated").send(emission)
|
||||
return emission.generation[0]
|
||||
|
||||
return narration_wrapper
|
||||
|
||||
|
||||
@register()
|
||||
class NarratorAgent(Agent):
|
||||
|
||||
"""
|
||||
Handles narration of the story
|
||||
"""
|
||||
|
||||
|
||||
agent_type = "narrator"
|
||||
verbose_name = "Narrator"
|
||||
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: client.TaleMateClient,
|
||||
**kwargs,
|
||||
):
|
||||
self.client = client
|
||||
|
||||
|
||||
# agent actions
|
||||
|
||||
|
||||
self.actions = {
|
||||
"narrate_time_passage": AgentAction(enabled=True, label="Narrate Time Passage", description="Whenever you indicate passage of time, narrate right after"),
|
||||
"generation_override": AgentAction(
|
||||
enabled=True,
|
||||
label="Generation Override",
|
||||
description="Override generation parameters",
|
||||
config={
|
||||
"instructions": AgentActionConfig(
|
||||
type="text",
|
||||
label="Instructions",
|
||||
value="Never wax poetic.",
|
||||
description="Extra instructions to give to the AI for narrative generation.",
|
||||
),
|
||||
},
|
||||
),
|
||||
"auto_break_repetition": AgentAction(
|
||||
enabled=True,
|
||||
label="Auto Break Repetition",
|
||||
description="Will attempt to automatically break AI repetition.",
|
||||
),
|
||||
"narrate_time_passage": AgentAction(
|
||||
enabled=True,
|
||||
label="Narrate Time Passage",
|
||||
description="Whenever you indicate passage of time, narrate right after",
|
||||
config={
|
||||
"ask_for_prompt": AgentActionConfig(
|
||||
type="bool",
|
||||
label="Guide time narration via prompt",
|
||||
description="Ask the user for a prompt to generate the time passage narration",
|
||||
value=True,
|
||||
)
|
||||
},
|
||||
),
|
||||
"narrate_dialogue": AgentAction(
|
||||
enabled=True,
|
||||
label="Narrate Dialogue",
|
||||
enabled=False,
|
||||
label="Narrate after Dialogue",
|
||||
description="Narrator will get a chance to narrate after every line of dialogue",
|
||||
config = {
|
||||
config={
|
||||
"ai_dialog": AgentActionConfig(
|
||||
type="number",
|
||||
label="AI Dialogue",
|
||||
label="AI Dialogue",
|
||||
description="Chance to narrate after every line of dialogue, 1 = always, 0 = never",
|
||||
value=0.3,
|
||||
value=0.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.1,
|
||||
),
|
||||
"player_dialog": AgentActionConfig(
|
||||
type="number",
|
||||
label="Player Dialogue",
|
||||
label="Player Dialogue",
|
||||
description="Chance to narrate after every line of dialogue, 1 = always, 0 = never",
|
||||
value=0.3,
|
||||
value=0.1,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.1,
|
||||
),
|
||||
}
|
||||
"generate_dialogue": AgentActionConfig(
|
||||
type="bool",
|
||||
label="Allow Dialogue in Narration",
|
||||
description="Allow the narrator to generate dialogue in narration",
|
||||
value=False,
|
||||
),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@property
|
||||
def extra_instructions(self):
|
||||
if self.actions["generation_override"].enabled:
|
||||
return self.actions["generation_override"].config["instructions"].value
|
||||
return ""
|
||||
|
||||
def clean_result(self, result):
|
||||
|
||||
"""
|
||||
Cleans the result of a narration
|
||||
"""
|
||||
|
||||
|
||||
result = result.strip().strip(":").strip()
|
||||
|
||||
|
||||
if "#" in result:
|
||||
result = result.split("#")[0]
|
||||
|
||||
|
||||
character_names = [c.name for c in self.scene.get_characters()]
|
||||
|
||||
|
||||
|
||||
cleaned = []
|
||||
for line in result.split("\n"):
|
||||
for character_name in character_names:
|
||||
@@ -90,58 +163,85 @@ class NarratorAgent(Agent):
|
||||
cleaned.append(line)
|
||||
|
||||
result = "\n".join(cleaned)
|
||||
#result = util.strip_partial_sentences(result)
|
||||
# result = util.strip_partial_sentences(result)
|
||||
return result
|
||||
|
||||
def connect(self, scene):
|
||||
|
||||
"""
|
||||
Connect to signals
|
||||
"""
|
||||
|
||||
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("agent.world_state.time").connect(self.on_time_passage)
|
||||
talemate.emit.async_signals.get("agent.world_state.time").connect(
|
||||
self.on_time_passage
|
||||
)
|
||||
talemate.emit.async_signals.get("game_loop_actor_iter").connect(self.on_dialog)
|
||||
|
||||
async def on_time_passage(self, event:TimePassageEmission):
|
||||
|
||||
async def on_time_passage(self, event: TimePassageEmission):
|
||||
"""
|
||||
Handles time passage narration, if enabled
|
||||
"""
|
||||
|
||||
|
||||
if not self.actions["narrate_time_passage"].enabled:
|
||||
return
|
||||
|
||||
response = await self.narrate_time_passage(event.duration, event.narrative)
|
||||
narrator_message = NarratorMessage(response, source=f"narrate_time_passage:{event.duration};{event.narrative}")
|
||||
|
||||
response = await self.narrate_time_passage(
|
||||
event.duration, event.human_duration, event.narrative
|
||||
)
|
||||
narrator_message = NarratorMessage(
|
||||
response, source=f"narrate_time_passage:{event.duration};{event.narrative}"
|
||||
)
|
||||
emit("narrator", narrator_message)
|
||||
self.scene.push_history(narrator_message)
|
||||
|
||||
async def on_dialog(self, event:GameLoopActorIterEvent):
|
||||
|
||||
|
||||
async def on_dialog(self, event: GameLoopActorIterEvent):
|
||||
"""
|
||||
Handles dialogue narration, if enabled
|
||||
"""
|
||||
|
||||
|
||||
if not self.actions["narrate_dialogue"].enabled:
|
||||
return
|
||||
|
||||
narrate_on_ai_chance = random.random() < self.actions["narrate_dialogue"].config["ai_dialog"].value
|
||||
narrate_on_player_chance = random.random() < self.actions["narrate_dialogue"].config["player_dialog"].value
|
||||
|
||||
log.debug("narrate on dialog", narrate_on_ai_chance=narrate_on_ai_chance, narrate_on_player_chance=narrate_on_player_chance)
|
||||
|
||||
if event.actor.character.is_player and not narrate_on_player_chance:
|
||||
|
||||
if event.game_loop.had_passive_narration:
|
||||
log.debug(
|
||||
"narrate on dialog",
|
||||
skip=True,
|
||||
had_passive_narration=event.game_loop.had_passive_narration,
|
||||
)
|
||||
return
|
||||
|
||||
if not event.actor.character.is_player and not narrate_on_ai_chance:
|
||||
|
||||
narrate_on_ai_chance = (
|
||||
self.actions["narrate_dialogue"].config["ai_dialog"].value
|
||||
)
|
||||
narrate_on_player_chance = (
|
||||
self.actions["narrate_dialogue"].config["player_dialog"].value
|
||||
)
|
||||
narrate_on_ai = random.random() < narrate_on_ai_chance
|
||||
narrate_on_player = random.random() < narrate_on_player_chance
|
||||
|
||||
log.debug(
|
||||
"narrate on dialog",
|
||||
narrate_on_ai=narrate_on_ai,
|
||||
narrate_on_ai_chance=narrate_on_ai_chance,
|
||||
narrate_on_player=narrate_on_player,
|
||||
narrate_on_player_chance=narrate_on_player_chance,
|
||||
)
|
||||
|
||||
if event.actor.character.is_player and not narrate_on_player:
|
||||
return
|
||||
|
||||
|
||||
if not event.actor.character.is_player and not narrate_on_ai:
|
||||
return
|
||||
|
||||
response = await self.narrate_after_dialogue(event.actor.character)
|
||||
narrator_message = NarratorMessage(response, source=f"narrate_dialogue:{event.actor.character.name}")
|
||||
narrator_message = NarratorMessage(
|
||||
response, source=f"narrate_dialogue:{event.actor.character.name}"
|
||||
)
|
||||
emit("narrator", narrator_message)
|
||||
self.scene.push_history(narrator_message)
|
||||
|
||||
event.game_loop.had_passive_narration = True
|
||||
|
||||
@set_processing
|
||||
async def narrate_scene(self):
|
||||
"""
|
||||
@@ -152,60 +252,51 @@ class NarratorAgent(Agent):
|
||||
"narrator.narrate-scene",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
response = response.strip("*")
|
||||
response = util.strip_partial_sentences(response)
|
||||
|
||||
|
||||
response = f"*{response.strip('*')}*"
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
async def progress_story(self, narrative_direction:str=None):
|
||||
async def progress_story(self, narrative_direction: str = None):
|
||||
"""
|
||||
Narrate the scene
|
||||
"""
|
||||
|
||||
scene = self.scene
|
||||
director = scene.get_helper("director").agent
|
||||
pc = scene.get_player_character()
|
||||
npcs = list(scene.get_npc_characters())
|
||||
npc_names= ", ".join([npc.name for npc in npcs])
|
||||
|
||||
#summarized_history = await scene.summarized_dialogue_history(
|
||||
# budget = self.client.max_token_length - 300,
|
||||
# min_dialogue = 50,
|
||||
#)
|
||||
|
||||
#augmented_context = await self.augment_context()
|
||||
npc_names = ", ".join([npc.name for npc in npcs])
|
||||
|
||||
if narrative_direction is None:
|
||||
#narrative_direction = await director.direct_narrative(
|
||||
# scene.context_history(budget=self.client.max_token_length - 500, min_dialogue=20),
|
||||
#)
|
||||
narrative_direction = "Slightly move the current scene forward."
|
||||
|
||||
self.scene.log.info("narrative_direction", narrative_direction=narrative_direction)
|
||||
|
||||
self.scene.log.info(
|
||||
"narrative_direction", narrative_direction=narrative_direction
|
||||
)
|
||||
|
||||
response = await Prompt.request(
|
||||
"narrator.narrate-progress",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
#"summarized_history": summarized_history,
|
||||
#"augmented_context": augmented_context,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"narrative_direction": narrative_direction,
|
||||
"player_character": pc,
|
||||
"npcs": npcs,
|
||||
"npc_names": npc_names,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
self.scene.log.info("progress_story", response=response)
|
||||
@@ -217,11 +308,13 @@ class NarratorAgent(Agent):
|
||||
if response.count("*") % 2 != 0:
|
||||
response = response.replace("*", "")
|
||||
response = f"*{response}*"
|
||||
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
async def narrate_query(self, query:str, at_the_end:bool=False, as_narrative:bool=True):
|
||||
async def narrate_query(
|
||||
self, query: str, at_the_end: bool = False, as_narrative: bool = True
|
||||
):
|
||||
"""
|
||||
Narrate a specific query
|
||||
"""
|
||||
@@ -229,20 +322,21 @@ class NarratorAgent(Agent):
|
||||
"narrator.narrate-query",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"query": query,
|
||||
"at_the_end": at_the_end,
|
||||
"as_narrative": as_narrative,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
log.info("narrate_query", response=response)
|
||||
response = self.clean_result(response.strip())
|
||||
log.info("narrate_query (after clean)", response=response)
|
||||
if as_narrative:
|
||||
response = f"*{response}*"
|
||||
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
@@ -251,27 +345,16 @@ class NarratorAgent(Agent):
|
||||
Narrate a specific character
|
||||
"""
|
||||
|
||||
budget = self.client.max_token_length - 300
|
||||
|
||||
memory_budget = min(int(budget * 0.05), 200)
|
||||
memory = self.scene.get_helper("memory").agent
|
||||
query = [
|
||||
f"What does {character.name} currently look like?",
|
||||
f"What is {character.name} currently wearing?",
|
||||
]
|
||||
memory_context = await memory.multi_query(
|
||||
query, iterate=1, max_tokens=memory_budget
|
||||
)
|
||||
response = await Prompt.request(
|
||||
"narrator.narrate-character",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"character": character,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"memory": memory_context,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
response = self.clean_result(response.strip())
|
||||
@@ -281,52 +364,55 @@ class NarratorAgent(Agent):
|
||||
|
||||
@set_processing
|
||||
async def augment_context(self):
|
||||
|
||||
"""
|
||||
Takes a context history generated via scene.context_history() and augments it with additional information
|
||||
by asking and answering questions with help from the long term memory.
|
||||
"""
|
||||
memory = self.scene.get_helper("memory").agent
|
||||
|
||||
|
||||
questions = await Prompt.request(
|
||||
"narrator.context-questions",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
self.scene.log.info("context_questions", questions=questions)
|
||||
|
||||
|
||||
questions = [q for q in questions.split("\n") if q.strip()]
|
||||
|
||||
|
||||
memory_context = await memory.multi_query(
|
||||
questions, iterate=2, max_tokens=self.client.max_token_length - 1000
|
||||
)
|
||||
|
||||
|
||||
answers = await Prompt.request(
|
||||
"narrator.context-answers",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"memory": memory_context,
|
||||
"questions": questions,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
self.scene.log.info("context_answers", answers=answers)
|
||||
|
||||
|
||||
answers = [a for a in answers.split("\n") if a.strip()]
|
||||
|
||||
|
||||
# return questions and answers
|
||||
return list(zip(questions, answers))
|
||||
|
||||
|
||||
@set_processing
|
||||
async def narrate_time_passage(self, duration:str, narrative:str=None):
|
||||
async def narrate_time_passage(
|
||||
self, duration: str, time_passed: str, narrative: str
|
||||
):
|
||||
"""
|
||||
Narrate a specific character
|
||||
"""
|
||||
@@ -335,24 +421,25 @@ class NarratorAgent(Agent):
|
||||
"narrator.narrate-time-passage",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"duration": duration,
|
||||
"time_passed": time_passed,
|
||||
"narrative": narrative,
|
||||
}
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
log.info("narrate_time_passage", response=response)
|
||||
|
||||
response = self.clean_result(response.strip())
|
||||
response = f"*{response}*"
|
||||
|
||||
return response
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def narrate_after_dialogue(self, character:Character):
|
||||
async def narrate_after_dialogue(self, character: Character):
|
||||
"""
|
||||
Narrate after a line of dialogue
|
||||
"""
|
||||
@@ -361,17 +448,196 @@ class NarratorAgent(Agent):
|
||||
"narrator.narrate-after-dialogue",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"character": character,
|
||||
"last_line": str(self.scene.history[-1])
|
||||
}
|
||||
"last_line": str(self.scene.history[-1]),
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
log.info("narrate_after_dialogue", response=response)
|
||||
|
||||
response = self.clean_result(response.strip().strip("*"))
|
||||
response = f"*{response}*"
|
||||
|
||||
return response
|
||||
allow_dialogue = (
|
||||
self.actions["narrate_dialogue"].config["generate_dialogue"].value
|
||||
)
|
||||
|
||||
if not allow_dialogue:
|
||||
response = response.split('"')[0].strip()
|
||||
response = response.replace("*", "")
|
||||
response = util.strip_partial_sentences(response)
|
||||
response = f"*{response}*"
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
async def narrate_character_entry(
|
||||
self, character: Character, direction: str = None
|
||||
):
|
||||
"""
|
||||
Narrate a character entering the scene
|
||||
"""
|
||||
|
||||
response = await Prompt.request(
|
||||
"narrator.narrate-character-entry",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"character": character,
|
||||
"direction": direction,
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
response = self.clean_result(response.strip().strip("*"))
|
||||
response = f"*{response}*"
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
async def narrate_character_exit(self, character: Character, direction: str = None):
|
||||
"""
|
||||
Narrate a character exiting the scene
|
||||
"""
|
||||
|
||||
response = await Prompt.request(
|
||||
"narrator.narrate-character-exit",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"character": character,
|
||||
"direction": direction,
|
||||
"extra_instructions": self.extra_instructions,
|
||||
},
|
||||
)
|
||||
|
||||
response = self.clean_result(response.strip().strip("*"))
|
||||
response = f"*{response}*"
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
async def paraphrase(self, narration: str):
|
||||
"""
|
||||
Paraphrase a narration
|
||||
"""
|
||||
|
||||
response = await Prompt.request(
|
||||
"narrator.paraphrase",
|
||||
self.client,
|
||||
"narrate",
|
||||
vars={
|
||||
"text": narration,
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
},
|
||||
)
|
||||
|
||||
log.info("paraphrase", narration=narration, response=response)
|
||||
|
||||
response = self.clean_result(response.strip().strip("*"))
|
||||
response = f"*{response}*"
|
||||
|
||||
return response
|
||||
|
||||
async def passthrough(self, narration: str) -> str:
|
||||
"""
|
||||
Pass through narration message as is
|
||||
"""
|
||||
narration = narration.replace("*", "")
|
||||
narration = f"*{narration}*"
|
||||
narration = util.ensure_dialog_format(narration)
|
||||
return narration
|
||||
|
||||
def action_to_source(
|
||||
self,
|
||||
action_name: str,
|
||||
parameters: dict,
|
||||
) -> str:
|
||||
"""
|
||||
Generate a source string for a given action and parameters
|
||||
|
||||
The source string is used to identify the source of a NarratorMessage
|
||||
and will also help regenerate the action and parameters from the source string
|
||||
later on
|
||||
"""
|
||||
|
||||
args = []
|
||||
|
||||
if action_name == "paraphrase":
|
||||
args.append(parameters.get("narration"))
|
||||
elif action_name == "narrate_character_entry":
|
||||
args.append(parameters.get("character").name)
|
||||
# args.append(parameters.get("direction"))
|
||||
elif action_name == "narrate_character_exit":
|
||||
args.append(parameters.get("character").name)
|
||||
# args.append(parameters.get("direction"))
|
||||
elif action_name == "narrate_character":
|
||||
args.append(parameters.get("character").name)
|
||||
elif action_name == "narrate_query":
|
||||
args.append(parameters.get("query"))
|
||||
elif action_name == "narrate_time_passage":
|
||||
args.append(parameters.get("duration"))
|
||||
args.append(parameters.get("time_passed"))
|
||||
args.append(parameters.get("narrative"))
|
||||
elif action_name == "progress_story":
|
||||
args.append(parameters.get("narrative_direction"))
|
||||
elif action_name == "narrate_after_dialogue":
|
||||
args.append(parameters.get("character"))
|
||||
|
||||
arg_str = ";".join(args) if args else ""
|
||||
|
||||
return f"{action_name}:{arg_str}".rstrip(":")
|
||||
|
||||
async def action_to_narration(
|
||||
self,
|
||||
action_name: str,
|
||||
emit_message: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
# calls self[action_name] and returns the result as a NarratorMessage
|
||||
# that is pushed to the history
|
||||
|
||||
fn = getattr(self, action_name)
|
||||
narration = await fn(**kwargs)
|
||||
source = self.action_to_source(action_name, kwargs)
|
||||
|
||||
narrator_message = NarratorMessage(narration, source=source)
|
||||
self.scene.push_history(narrator_message)
|
||||
|
||||
if emit_message:
|
||||
emit("narrator", narrator_message)
|
||||
|
||||
return narrator_message
|
||||
|
||||
# LLM client related methods. These are called during or after the client
|
||||
|
||||
def inject_prompt_paramters(
|
||||
self, prompt_param: dict, kind: str, agent_function_name: str
|
||||
):
|
||||
log.debug(
|
||||
"inject_prompt_paramters",
|
||||
prompt_param=prompt_param,
|
||||
kind=kind,
|
||||
agent_function_name=agent_function_name,
|
||||
)
|
||||
character_names = [f"\n{c.name}:" for c in self.scene.get_characters()]
|
||||
if prompt_param.get("extra_stopping_strings") is None:
|
||||
prompt_param["extra_stopping_strings"] = []
|
||||
prompt_param["extra_stopping_strings"] += character_names
|
||||
|
||||
def allow_repetition_break(
|
||||
self, kind: str, agent_function_name: str, auto: bool = False
|
||||
):
|
||||
if auto and not self.actions["auto_break_repetition"].enabled:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@@ -1,24 +1,26 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
import traceback
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import structlog
|
||||
|
||||
import talemate.data_objects as data_objects
|
||||
import talemate.emit.async_signals
|
||||
import talemate.util as util
|
||||
from talemate.events import GameLoopEvent
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import DirectorMessage, TimePassageMessage
|
||||
|
||||
from .base import Agent, set_processing
|
||||
from .base import Agent, AgentAction, AgentActionConfig, set_processing
|
||||
from .registry import register
|
||||
|
||||
import structlog
|
||||
|
||||
import time
|
||||
import re
|
||||
|
||||
log = structlog.get_logger("talemate.agents.summarize")
|
||||
|
||||
|
||||
@register()
|
||||
class SummarizeAgent(Agent):
|
||||
"""
|
||||
@@ -35,12 +37,56 @@ class SummarizeAgent(Agent):
|
||||
def __init__(self, client, **kwargs):
|
||||
self.client = client
|
||||
|
||||
def on_history_add(self, event):
|
||||
asyncio.ensure_future(self.build_archive(event.scene))
|
||||
self.actions = {
|
||||
"archive": AgentAction(
|
||||
enabled=True,
|
||||
label="Summarize to long-term memory archive",
|
||||
description="Automatically summarize scene dialogue when the number of tokens in the history exceeds a threshold. This helps keep the context history from growing too large.",
|
||||
config={
|
||||
"threshold": AgentActionConfig(
|
||||
type="number",
|
||||
label="Token Threshold",
|
||||
description="Will summarize when the number of tokens in the history exceeds this threshold",
|
||||
min=512,
|
||||
max=8192,
|
||||
step=256,
|
||||
value=1536,
|
||||
),
|
||||
"method": AgentActionConfig(
|
||||
type="text",
|
||||
label="Summarization Method",
|
||||
description="Which method to use for summarization",
|
||||
value="balanced",
|
||||
choices=[
|
||||
{"label": "Short & Concise", "value": "short"},
|
||||
{"label": "Balanced", "value": "balanced"},
|
||||
{"label": "Lengthy & Detailed", "value": "long"},
|
||||
],
|
||||
),
|
||||
"include_previous": AgentActionConfig(
|
||||
type="number",
|
||||
label="Use preceeding summaries to strengthen context",
|
||||
description="Number of entries",
|
||||
note="Help the AI summarize by including the last few summaries as additional context. Some models may incorporate this context into the new summary directly, so if you find yourself with a bunch of similar history entries, try setting this to 0.",
|
||||
value=3,
|
||||
min=0,
|
||||
max=10,
|
||||
step=1,
|
||||
),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
def connect(self, scene):
|
||||
super().connect(scene)
|
||||
scene.signals["history_add"].connect(self.on_history_add)
|
||||
talemate.emit.async_signals.get("game_loop").connect(self.on_game_loop)
|
||||
|
||||
async def on_game_loop(self, emission: GameLoopEvent):
|
||||
"""
|
||||
Called when a conversation is generated
|
||||
"""
|
||||
|
||||
await self.build_archive(self.scene)
|
||||
|
||||
def clean_result(self, result):
|
||||
if "#" in result:
|
||||
@@ -53,45 +99,75 @@ class SummarizeAgent(Agent):
|
||||
return result
|
||||
|
||||
@set_processing
|
||||
async def build_archive(self, scene, token_threshold:int=1500):
|
||||
async def build_archive(self, scene):
|
||||
end = None
|
||||
|
||||
if not self.actions["archive"].enabled:
|
||||
return
|
||||
|
||||
if not scene.archived_history:
|
||||
start = 0
|
||||
recent_entry = None
|
||||
else:
|
||||
recent_entry = scene.archived_history[-1]
|
||||
start = recent_entry.get("end", 0) + 1
|
||||
if "end" not in recent_entry:
|
||||
# permanent historical archive entry, not tied to any specific history entry
|
||||
# meaning we are still at the beginning of the scene
|
||||
start = 0
|
||||
else:
|
||||
start = recent_entry.get("end", 0) + 1
|
||||
|
||||
# if there is a recent entry we also collect the 3 most recentries
|
||||
# as extra context
|
||||
|
||||
num_previous = self.actions["archive"].config["include_previous"].value
|
||||
if recent_entry and num_previous > 0:
|
||||
extra_context = "\n\n".join(
|
||||
[entry["text"] for entry in scene.archived_history[-num_previous:]]
|
||||
)
|
||||
else:
|
||||
extra_context = None
|
||||
|
||||
tokens = 0
|
||||
dialogue_entries = []
|
||||
ts = "PT0S"
|
||||
time_passage_termination = False
|
||||
|
||||
|
||||
token_threshold = self.actions["archive"].config["threshold"].value
|
||||
|
||||
log.debug("build_archive", start=start, recent_entry=recent_entry)
|
||||
|
||||
|
||||
if recent_entry:
|
||||
ts = recent_entry.get("ts", ts)
|
||||
|
||||
|
||||
for i in range(start, len(scene.history)):
|
||||
dialogue = scene.history[i]
|
||||
|
||||
# log.debug("build_archive", idx=i, content=str(dialogue)[:64]+"...")
|
||||
|
||||
if isinstance(dialogue, DirectorMessage):
|
||||
if i == start:
|
||||
start += 1
|
||||
continue
|
||||
|
||||
|
||||
if isinstance(dialogue, TimePassageMessage):
|
||||
log.debug("build_archive", time_passage_message=dialogue)
|
||||
if i == start:
|
||||
ts = util.iso8601_add(ts, dialogue.ts)
|
||||
log.debug("build_archive", time_passage_message=dialogue, start=start, i=i, ts=ts)
|
||||
log.debug(
|
||||
"build_archive",
|
||||
time_passage_message=dialogue,
|
||||
start=start,
|
||||
i=i,
|
||||
ts=ts,
|
||||
)
|
||||
start += 1
|
||||
continue
|
||||
log.debug("build_archive", time_passage_message_termination=dialogue)
|
||||
time_passage_termination = True
|
||||
end = i - 1
|
||||
break
|
||||
|
||||
|
||||
tokens += util.count_tokens(dialogue)
|
||||
dialogue_entries.append(dialogue)
|
||||
if tokens > token_threshold: #
|
||||
@@ -101,43 +177,44 @@ class SummarizeAgent(Agent):
|
||||
if end is None:
|
||||
# nothing to archive yet
|
||||
return
|
||||
|
||||
log.debug("build_archive", start=start, end=end, ts=ts, time_passage_termination=time_passage_termination)
|
||||
|
||||
extra_context = None
|
||||
if recent_entry:
|
||||
extra_context = recent_entry["text"]
|
||||
|
||||
log.debug(
|
||||
"build_archive",
|
||||
start=start,
|
||||
end=end,
|
||||
ts=ts,
|
||||
time_passage_termination=time_passage_termination,
|
||||
)
|
||||
|
||||
# in order to summarize coherently, we need to determine if there is a favorable
|
||||
# cutoff point (e.g., the scene naturally ends or shifts meaninfully in the middle
|
||||
# of the dialogue)
|
||||
#
|
||||
# One way to do this is to check if the last line is a TimePassageMessage, which
|
||||
# indicates a scene change or a significant pause.
|
||||
#
|
||||
# indicates a scene change or a significant pause.
|
||||
#
|
||||
# If not, we can ask the AI to find a good point of
|
||||
# termination.
|
||||
|
||||
|
||||
if not time_passage_termination:
|
||||
|
||||
# No TimePassageMessage, so we need to ask the AI to find a good point of termination
|
||||
|
||||
|
||||
terminating_line = await self.analyze_dialoge(dialogue_entries)
|
||||
|
||||
if terminating_line:
|
||||
adjusted_dialogue = []
|
||||
for line in dialogue_entries:
|
||||
for line in dialogue_entries:
|
||||
if str(line) in terminating_line:
|
||||
break
|
||||
adjusted_dialogue.append(line)
|
||||
dialogue_entries = adjusted_dialogue
|
||||
end = start + len(dialogue_entries)
|
||||
|
||||
end = start + len(dialogue_entries) - 1
|
||||
|
||||
if dialogue_entries:
|
||||
summarized = await self.summarize(
|
||||
"\n".join(map(str, dialogue_entries)), extra_context=extra_context
|
||||
)
|
||||
|
||||
|
||||
else:
|
||||
# AI has likely identified the first line as a scene change, so we can't summarize
|
||||
# just use the first line
|
||||
@@ -151,55 +228,178 @@ class SummarizeAgent(Agent):
|
||||
|
||||
@set_processing
|
||||
async def analyze_dialoge(self, dialogue):
|
||||
response = await Prompt.request("summarizer.analyze-dialogue", self.client, "analyze_freeform", vars={
|
||||
"dialogue": "\n".join(map(str, dialogue)),
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
})
|
||||
|
||||
response = await Prompt.request(
|
||||
"summarizer.analyze-dialogue",
|
||||
self.client,
|
||||
"analyze_freeform",
|
||||
vars={
|
||||
"dialogue": "\n".join(map(str, dialogue)),
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
},
|
||||
)
|
||||
|
||||
response = self.clean_result(response)
|
||||
return response
|
||||
|
||||
|
||||
@set_processing
|
||||
async def summarize(
|
||||
self,
|
||||
text: str,
|
||||
perspective: str = None,
|
||||
pins: Union[List[str], None] = None,
|
||||
extra_context: str = None,
|
||||
method: str = None,
|
||||
extra_instructions: str = None,
|
||||
):
|
||||
"""
|
||||
Summarize the given text
|
||||
"""
|
||||
|
||||
response = await Prompt.request("summarizer.summarize-dialogue", self.client, "summarize", vars={
|
||||
"dialogue": text,
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
})
|
||||
|
||||
self.scene.log.info("summarize", dialogue_length=len(text), summarized_length=len(response))
|
||||
response = await Prompt.request(
|
||||
"summarizer.summarize-dialogue",
|
||||
self.client,
|
||||
"summarize",
|
||||
vars={
|
||||
"dialogue": text,
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"summarization_method": (
|
||||
self.actions["archive"].config["method"].value
|
||||
if method is None
|
||||
else method
|
||||
),
|
||||
"extra_context": extra_context or "",
|
||||
"extra_instructions": extra_instructions or "",
|
||||
},
|
||||
)
|
||||
|
||||
self.scene.log.info(
|
||||
"summarize", dialogue_length=len(text), summarized_length=len(response)
|
||||
)
|
||||
|
||||
return self.clean_result(response)
|
||||
|
||||
@set_processing
|
||||
async def simple_summary(
|
||||
self, text: str, prompt_kind: str = "summarize", instructions: str = "Summarize"
|
||||
):
|
||||
prompt = [
|
||||
text,
|
||||
"",
|
||||
f"Instruction: {instructions}",
|
||||
"<|BOT|>Short Summary: ",
|
||||
async def build_stepped_archive_for_level(self, level: int):
|
||||
"""
|
||||
WIP - not yet used
|
||||
|
||||
This will iterate over existing archived_history entries
|
||||
and stepped_archived_history entries and summarize based on time duration
|
||||
indicated between the entries.
|
||||
|
||||
The lowest level of summarization (based on token threshold and any time passage)
|
||||
happens in build_archive. This method is for summarizing furhter levels based on
|
||||
long time pasages.
|
||||
|
||||
Level 0: small timestap summarize (summarizes all token summarizations when time advances +1 day)
|
||||
Level 1: medium timestap summarize (summarizes all small timestep summarizations when time advances +1 week)
|
||||
Level 2: large timestap summarize (summarizes all medium timestep summarizations when time advances +1 month)
|
||||
Level 3: huge timestap summarize (summarizes all large timestep summarizations when time advances +1 year)
|
||||
Level 4: massive timestap summarize (summarizes all huge timestep summarizations when time advances +10 years)
|
||||
Level 5: epic timestap summarize (summarizes all massive timestep summarizations when time advances +100 years)
|
||||
and so on (increasing by a factor of 10 each time)
|
||||
|
||||
```
|
||||
@dataclass
|
||||
class ArchiveEntry:
|
||||
text: str
|
||||
start: int = None
|
||||
end: int = None
|
||||
ts: str = None
|
||||
```
|
||||
|
||||
Like token summarization this will use ArchiveEntry and start and end will refer to the entries in the
|
||||
lower level of summarization.
|
||||
|
||||
Ts is the iso8601 timestamp of the start of the summarized period.
|
||||
"""
|
||||
|
||||
# select the list to use for the entries
|
||||
|
||||
if level == 0:
|
||||
entries = self.scene.archived_history
|
||||
else:
|
||||
entries = self.scene.stepped_archived_history[level - 1]
|
||||
|
||||
# select the list to summarize new entries to
|
||||
|
||||
target = self.scene.stepped_archived_history[level]
|
||||
|
||||
if not target:
|
||||
raise ValueError(f"Invalid level {level}")
|
||||
|
||||
# determine the start and end of the period to summarize
|
||||
|
||||
if not entries:
|
||||
return
|
||||
|
||||
# determine the time threshold for this level
|
||||
|
||||
# first calculate all possible thresholds in iso8601 format, starting with 1 day
|
||||
thresholds = [
|
||||
"P1D",
|
||||
"P1W",
|
||||
"P1M",
|
||||
"P1Y",
|
||||
]
|
||||
|
||||
response = await self.client.send_prompt("\n".join(map(str, prompt)), kind=prompt_kind)
|
||||
if ":" in response:
|
||||
response = response.split(":")[1].strip()
|
||||
return response
|
||||
# TODO: auto extend?
|
||||
|
||||
time_threshold_in_seconds = util.iso8601_to_seconds(thresholds[level])
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
if not time_threshold_in_seconds:
|
||||
raise ValueError(f"Invalid level {level}")
|
||||
|
||||
# determine the most recent summarized entry time, and then find entries
|
||||
# that are newer than that in the lower list
|
||||
|
||||
ts = target[-1].ts if target else entries[0].ts
|
||||
|
||||
# determine the most recent entry at the lower level, if its not newer or
|
||||
# the difference is less than the threshold, then we don't need to summarize
|
||||
|
||||
recent_entry = entries[-1]
|
||||
|
||||
if util.iso8601_diff(recent_entry.ts, ts) < time_threshold_in_seconds:
|
||||
return
|
||||
|
||||
log.debug("build_stepped_archive", level=level, ts=ts)
|
||||
|
||||
# if target is empty, start is 0
|
||||
# otherwise start is the end of the last entry
|
||||
|
||||
start = 0 if not target else target[-1].end
|
||||
|
||||
# collect entries starting at start until the combined time duration
|
||||
# exceeds the threshold
|
||||
|
||||
entries_to_summarize = []
|
||||
|
||||
for entry in entries[start:]:
|
||||
entries_to_summarize.append(entry)
|
||||
if util.iso8601_diff(entry.ts, ts) > time_threshold_in_seconds:
|
||||
break
|
||||
|
||||
# summarize the entries
|
||||
# we also collect N entries of previous summaries to use as context
|
||||
|
||||
num_previous = self.actions["archive"].config["include_previous"].value
|
||||
if num_previous > 0:
|
||||
extra_context = "\n\n".join(
|
||||
[entry["text"] for entry in target[-num_previous:]]
|
||||
)
|
||||
else:
|
||||
extra_context = None
|
||||
|
||||
summarized = await self.summarize(
|
||||
"\n".join(map(str, entries_to_summarize)), extra_context=extra_context
|
||||
)
|
||||
|
||||
# push summarized entry to target
|
||||
|
||||
ts = entries_to_summarize[-1].ts
|
||||
|
||||
target.append(
|
||||
data_objects.ArchiveEntry(
|
||||
summarized, start, len(entries_to_summarize) - 1, ts=ts
|
||||
)
|
||||
)
|
||||
|
||||
652
src/talemate/agents/tts.py
Normal file
@@ -0,0 +1,652 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import functools
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from typing import Union
|
||||
|
||||
import httpx
|
||||
import nltk
|
||||
import pydantic
|
||||
import structlog
|
||||
from nltk.tokenize import sent_tokenize
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
import talemate.config as config
|
||||
import talemate.emit.async_signals
|
||||
import talemate.instance as instance
|
||||
from talemate.emit import emit
|
||||
from talemate.emit.signals import handlers
|
||||
from talemate.events import GameLoopNewMessageEvent
|
||||
from talemate.scene_message import CharacterMessage, NarratorMessage
|
||||
|
||||
from .base import (
|
||||
Agent,
|
||||
AgentAction,
|
||||
AgentActionConditional,
|
||||
AgentActionConfig,
|
||||
AgentDetail,
|
||||
set_processing,
|
||||
)
|
||||
from .registry import register
|
||||
|
||||
try:
|
||||
from TTS.api import TTS
|
||||
except ImportError:
|
||||
TTS = None
|
||||
|
||||
log = structlog.get_logger("talemate.agents.tts") #
|
||||
|
||||
if not TTS:
|
||||
# TTS installation is massive and requires a lot of dependencies
|
||||
# so we don't want to require it unless the user wants to use it
|
||||
log.info(
|
||||
"TTS (local) requires the TTS package, please install with `pip install TTS` if you want to use the local api"
|
||||
)
|
||||
|
||||
|
||||
def parse_chunks(text):
|
||||
text = text.replace("...", "__ellipsis__")
|
||||
|
||||
chunks = sent_tokenize(text)
|
||||
cleaned_chunks = []
|
||||
|
||||
for chunk in chunks:
|
||||
chunk = chunk.replace("*", "")
|
||||
if not chunk:
|
||||
continue
|
||||
cleaned_chunks.append(chunk)
|
||||
|
||||
for i, chunk in enumerate(cleaned_chunks):
|
||||
chunk = chunk.replace("__ellipsis__", "...")
|
||||
|
||||
cleaned_chunks[i] = chunk
|
||||
|
||||
return cleaned_chunks
|
||||
|
||||
|
||||
def clean_quotes(chunk: str):
|
||||
# if there is an uneven number of quotes, remove the last one if its
|
||||
# at the end of the chunk. If its in the middle, add a quote to the end
|
||||
if chunk.count('"') % 2 == 1:
|
||||
if chunk.endswith('"'):
|
||||
chunk = chunk[:-1]
|
||||
else:
|
||||
chunk += '"'
|
||||
|
||||
return chunk
|
||||
|
||||
|
||||
def rejoin_chunks(chunks: list[str], chunk_size: int = 250):
|
||||
"""
|
||||
Will combine chunks split by punctuation into a single chunk until
|
||||
max chunk size is reached
|
||||
"""
|
||||
|
||||
joined_chunks = []
|
||||
|
||||
current_chunk = ""
|
||||
|
||||
for chunk in chunks:
|
||||
if len(current_chunk) + len(chunk) > chunk_size:
|
||||
joined_chunks.append(clean_quotes(current_chunk))
|
||||
current_chunk = ""
|
||||
|
||||
current_chunk += chunk
|
||||
|
||||
if current_chunk:
|
||||
joined_chunks.append(clean_quotes(current_chunk))
|
||||
return joined_chunks
|
||||
|
||||
|
||||
class Voice(pydantic.BaseModel):
|
||||
value: str
|
||||
label: str
|
||||
|
||||
|
||||
class VoiceLibrary(pydantic.BaseModel):
|
||||
api: str
|
||||
voices: list[Voice] = pydantic.Field(default_factory=list)
|
||||
last_synced: float = None
|
||||
|
||||
|
||||
@register()
|
||||
class TTSAgent(Agent):
|
||||
"""
|
||||
Text to speech agent
|
||||
"""
|
||||
|
||||
agent_type = "tts"
|
||||
verbose_name = "Voice"
|
||||
requires_llm_client = False
|
||||
essential = False
|
||||
|
||||
@classmethod
|
||||
def config_options(cls, agent=None):
|
||||
config_options = super().config_options(agent=agent)
|
||||
|
||||
if agent:
|
||||
config_options["actions"]["_config"]["config"]["voice_id"]["choices"] = [
|
||||
voice.model_dump() for voice in agent.list_voices_sync()
|
||||
]
|
||||
|
||||
return config_options
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.is_enabled = False
|
||||
nltk.download("punkt", quiet=True)
|
||||
|
||||
self.voices = {
|
||||
"elevenlabs": VoiceLibrary(api="elevenlabs"),
|
||||
"tts": VoiceLibrary(api="tts"),
|
||||
"openai": VoiceLibrary(api="openai"),
|
||||
}
|
||||
self.config = config.load_config()
|
||||
self.playback_done_event = asyncio.Event()
|
||||
self.preselect_voice = None
|
||||
self.actions = {
|
||||
"_config": AgentAction(
|
||||
enabled=True,
|
||||
label="Configure",
|
||||
description="TTS agent configuration",
|
||||
config={
|
||||
"api": AgentActionConfig(
|
||||
type="text",
|
||||
choices=[
|
||||
{"value": "tts", "label": "TTS (Local)"},
|
||||
{"value": "elevenlabs", "label": "Eleven Labs"},
|
||||
{"value": "openai", "label": "OpenAI"},
|
||||
],
|
||||
value="tts",
|
||||
label="API",
|
||||
description="Which TTS API to use",
|
||||
onchange="emit",
|
||||
),
|
||||
"voice_id": AgentActionConfig(
|
||||
type="text",
|
||||
value="default",
|
||||
label="Narrator Voice",
|
||||
description="Voice ID/Name to use for TTS",
|
||||
choices=[],
|
||||
),
|
||||
"generate_for_player": AgentActionConfig(
|
||||
type="bool",
|
||||
value=False,
|
||||
label="Generate for player",
|
||||
description="Generate audio for player messages",
|
||||
),
|
||||
"generate_for_npc": AgentActionConfig(
|
||||
type="bool",
|
||||
value=True,
|
||||
label="Generate for NPCs",
|
||||
description="Generate audio for NPC messages",
|
||||
),
|
||||
"generate_for_narration": AgentActionConfig(
|
||||
type="bool",
|
||||
value=True,
|
||||
label="Generate for narration",
|
||||
description="Generate audio for narration messages",
|
||||
),
|
||||
"generate_chunks": AgentActionConfig(
|
||||
type="bool",
|
||||
value=False,
|
||||
label="Split generation",
|
||||
description="Generate audio chunks for each sentence - will be much more responsive but may loose context to inform inflection",
|
||||
),
|
||||
},
|
||||
),
|
||||
"openai": AgentAction(
|
||||
enabled=True,
|
||||
condition=AgentActionConditional(
|
||||
attribute="_config.config.api", value="openai"
|
||||
),
|
||||
label="OpenAI Settings",
|
||||
config={
|
||||
"model": AgentActionConfig(
|
||||
type="text",
|
||||
value="tts-1",
|
||||
choices=[
|
||||
{"value": "tts-1", "label": "TTS 1"},
|
||||
{"value": "tts-1-hd", "label": "TTS 1 HD"},
|
||||
],
|
||||
label="Model",
|
||||
description="TTS model to use",
|
||||
),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
self.actions["_config"].model_dump()
|
||||
handlers["config_saved"].connect(self.on_config_saved)
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return self.is_enabled
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
return True
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return False
|
||||
|
||||
@property
|
||||
def not_ready_reason(self) -> str:
|
||||
"""
|
||||
Returns a string explaining why the agent is not ready
|
||||
"""
|
||||
|
||||
if self.ready:
|
||||
return ""
|
||||
|
||||
if self.api == "tts":
|
||||
if not TTS:
|
||||
return "TTS not installed"
|
||||
|
||||
elif self.requires_token and not self.token:
|
||||
return "No API token"
|
||||
|
||||
elif not self.default_voice_id:
|
||||
return "No voice selected"
|
||||
|
||||
@property
|
||||
def agent_details(self):
|
||||
|
||||
details = {
|
||||
"api": AgentDetail(
|
||||
icon="mdi-server-outline",
|
||||
value=self.api_label,
|
||||
description="The backend to use for TTS",
|
||||
).model_dump(),
|
||||
}
|
||||
|
||||
if self.ready and self.enabled:
|
||||
details["voice"] = AgentDetail(
|
||||
icon="mdi-account-voice",
|
||||
value=self.voice_id_to_label(self.default_voice_id) or "",
|
||||
description="The voice to use for TTS",
|
||||
color="info",
|
||||
).model_dump()
|
||||
elif self.enabled:
|
||||
details["error"] = AgentDetail(
|
||||
icon="mdi-alert",
|
||||
value=self.not_ready_reason,
|
||||
description=self.not_ready_reason,
|
||||
color="error",
|
||||
).model_dump()
|
||||
|
||||
return details
|
||||
|
||||
@property
|
||||
def api(self):
|
||||
return self.actions["_config"].config["api"].value
|
||||
|
||||
@property
|
||||
def api_label(self):
|
||||
choices = self.actions["_config"].config["api"].choices
|
||||
api = self.api
|
||||
for choice in choices:
|
||||
if choice["value"] == api:
|
||||
return choice["label"]
|
||||
return api
|
||||
|
||||
@property
|
||||
def token(self):
|
||||
api = self.api
|
||||
return self.config.get(api, {}).get("api_key")
|
||||
|
||||
@property
|
||||
def default_voice_id(self):
|
||||
return self.actions["_config"].config["voice_id"].value
|
||||
|
||||
@property
|
||||
def requires_token(self):
|
||||
return self.api != "tts"
|
||||
|
||||
@property
|
||||
def ready(self):
|
||||
if self.api == "tts":
|
||||
if not TTS:
|
||||
return False
|
||||
return True
|
||||
|
||||
return (not self.requires_token or self.token) and self.default_voice_id
|
||||
|
||||
@property
|
||||
def status(self):
|
||||
if not self.enabled:
|
||||
return "disabled"
|
||||
if self.ready:
|
||||
if getattr(self, "processing_bg", 0) > 0:
|
||||
return "busy_bg" if not getattr(self, "processing", False) else "busy"
|
||||
return "active" if not getattr(self, "processing", False) else "busy"
|
||||
if self.requires_token and not self.token:
|
||||
return "error"
|
||||
if self.api == "tts":
|
||||
if not TTS:
|
||||
return "error"
|
||||
return "uninitialized"
|
||||
|
||||
@property
|
||||
def max_generation_length(self):
|
||||
if self.api == "elevenlabs":
|
||||
return 1024
|
||||
elif self.api == "coqui":
|
||||
return 250
|
||||
|
||||
return 250
|
||||
|
||||
@property
|
||||
def openai_api_key(self):
|
||||
return self.config.get("openai", {}).get("api_key")
|
||||
|
||||
async def apply_config(self, *args, **kwargs):
|
||||
try:
|
||||
api = kwargs["actions"]["_config"]["config"]["api"]["value"]
|
||||
except KeyError:
|
||||
api = self.api
|
||||
|
||||
api_changed = api != self.api
|
||||
|
||||
log.debug(
|
||||
"apply_config",
|
||||
api=api,
|
||||
api_changed=api != self.api,
|
||||
current_api=self.api,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
try:
|
||||
self.preselect_voice = kwargs["actions"]["_config"]["config"]["voice_id"][
|
||||
"value"
|
||||
]
|
||||
except KeyError:
|
||||
self.preselect_voice = self.default_voice_id
|
||||
|
||||
await super().apply_config(*args, **kwargs)
|
||||
|
||||
if api_changed:
|
||||
try:
|
||||
self.actions["_config"].config["voice_id"].value = (
|
||||
self.voices[api].voices[0].value
|
||||
)
|
||||
except IndexError:
|
||||
self.actions["_config"].config["voice_id"].value = ""
|
||||
|
||||
def connect(self, scene):
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("game_loop_new_message").connect(
|
||||
self.on_game_loop_new_message
|
||||
)
|
||||
|
||||
def on_config_saved(self, event):
|
||||
config = event.data
|
||||
self.config = config
|
||||
instance.emit_agent_status(self.__class__, self)
|
||||
|
||||
async def on_game_loop_new_message(self, emission: GameLoopNewMessageEvent):
|
||||
"""
|
||||
Called when a conversation is generated
|
||||
"""
|
||||
|
||||
if not self.enabled or not self.ready:
|
||||
return
|
||||
|
||||
if not isinstance(emission.message, (CharacterMessage, NarratorMessage)):
|
||||
return
|
||||
|
||||
if (
|
||||
isinstance(emission.message, NarratorMessage)
|
||||
and not self.actions["_config"].config["generate_for_narration"].value
|
||||
):
|
||||
return
|
||||
|
||||
if isinstance(emission.message, CharacterMessage):
|
||||
if (
|
||||
emission.message.source == "player"
|
||||
and not self.actions["_config"].config["generate_for_player"].value
|
||||
):
|
||||
return
|
||||
elif (
|
||||
emission.message.source == "ai"
|
||||
and not self.actions["_config"].config["generate_for_npc"].value
|
||||
):
|
||||
return
|
||||
|
||||
if isinstance(emission.message, CharacterMessage):
|
||||
character_prefix = emission.message.split(":", 1)[0]
|
||||
else:
|
||||
character_prefix = ""
|
||||
|
||||
log.info(
|
||||
"reactive tts", message=emission.message, character_prefix=character_prefix
|
||||
)
|
||||
|
||||
await self.generate(str(emission.message).replace(character_prefix + ": ", ""))
|
||||
|
||||
def voice(self, voice_id: str) -> Union[Voice, None]:
|
||||
for voice in self.voices[self.api].voices:
|
||||
if voice.value == voice_id:
|
||||
return voice
|
||||
return None
|
||||
|
||||
def voice_id_to_label(self, voice_id: str):
|
||||
for voice in self.voices[self.api].voices:
|
||||
if voice.value == voice_id:
|
||||
return voice.label
|
||||
return None
|
||||
|
||||
def list_voices_sync(self):
|
||||
loop = asyncio.get_event_loop()
|
||||
return loop.run_until_complete(self.list_voices())
|
||||
|
||||
async def list_voices(self):
|
||||
if self.requires_token and not self.token:
|
||||
return []
|
||||
|
||||
library = self.voices[self.api]
|
||||
|
||||
# TODO: allow re-syncing voices
|
||||
if library.last_synced:
|
||||
return library.voices
|
||||
|
||||
list_fn = getattr(self, f"_list_voices_{self.api}")
|
||||
log.info("Listing voices", api=self.api)
|
||||
|
||||
library.voices = await list_fn()
|
||||
library.last_synced = time.time()
|
||||
|
||||
if self.preselect_voice:
|
||||
if self.voice(self.preselect_voice):
|
||||
self.actions["_config"].config["voice_id"].value = self.preselect_voice
|
||||
self.preselect_voice = None
|
||||
|
||||
# if the current voice cannot be found, reset it
|
||||
if not self.voice(self.default_voice_id):
|
||||
self.actions["_config"].config["voice_id"].value = ""
|
||||
|
||||
# set loading to false
|
||||
return library.voices
|
||||
|
||||
@set_processing
|
||||
async def generate(self, text: str):
|
||||
if not self.enabled or not self.ready or not text:
|
||||
return
|
||||
|
||||
self.playback_done_event.set()
|
||||
|
||||
generate_fn = getattr(self, f"_generate_{self.api}")
|
||||
|
||||
if self.actions["_config"].config["generate_chunks"].value:
|
||||
chunks = parse_chunks(text)
|
||||
chunks = rejoin_chunks(chunks)
|
||||
else:
|
||||
chunks = parse_chunks(text)
|
||||
chunks = rejoin_chunks(chunks, chunk_size=self.max_generation_length)
|
||||
|
||||
# Start generating audio chunks in the background
|
||||
generation_task = asyncio.create_task(self.generate_chunks(generate_fn, chunks))
|
||||
await self.set_background_processing(generation_task)
|
||||
|
||||
# Wait for both tasks to complete
|
||||
# await asyncio.gather(generation_task)
|
||||
|
||||
async def generate_chunks(self, generate_fn, chunks):
|
||||
for chunk in chunks:
|
||||
chunk = chunk.replace("*", "").strip()
|
||||
log.info("Generating audio", api=self.api, chunk=chunk)
|
||||
audio_data = await generate_fn(chunk)
|
||||
self.play_audio(audio_data)
|
||||
|
||||
def play_audio(self, audio_data):
|
||||
# play audio through the python audio player
|
||||
# play(audio_data)
|
||||
|
||||
emit(
|
||||
"audio_queue",
|
||||
data={"audio_data": base64.b64encode(audio_data).decode("utf-8")},
|
||||
)
|
||||
|
||||
self.playback_done_event.set() # Signal that playback is finished
|
||||
|
||||
# LOCAL
|
||||
|
||||
async def _generate_tts(self, text: str) -> Union[bytes, None]:
|
||||
if not TTS:
|
||||
return
|
||||
|
||||
tts_config = self.config.get("tts", {})
|
||||
model = tts_config.get("model")
|
||||
device = tts_config.get("device", "cpu")
|
||||
|
||||
log.debug("tts local", model=model, device=device)
|
||||
|
||||
if not hasattr(self, "tts_instance"):
|
||||
self.tts_instance = TTS(model).to(device)
|
||||
|
||||
tts = self.tts_instance
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
voice = self.voice(self.default_voice_id)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, f"tts-{uuid.uuid4()}.wav")
|
||||
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
functools.partial(
|
||||
tts.tts_to_file,
|
||||
text=text,
|
||||
speaker_wav=voice.value,
|
||||
language="en",
|
||||
file_path=file_path,
|
||||
),
|
||||
)
|
||||
# tts.tts_to_file(text=text, speaker_wav=voice.value, language="en", file_path=file_path)
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
async def _list_voices_tts(self) -> dict[str, str]:
|
||||
return [
|
||||
Voice(**voice) for voice in self.config.get("tts", {}).get("voices", [])
|
||||
]
|
||||
|
||||
# ELEVENLABS
|
||||
|
||||
async def _generate_elevenlabs(
|
||||
self, text: str, chunk_size: int = 1024
|
||||
) -> Union[bytes, None]:
|
||||
api_key = self.token
|
||||
if not api_key:
|
||||
return
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
url = f"https://api.elevenlabs.io/v1/text-to-speech/{self.default_voice_id}"
|
||||
headers = {
|
||||
"Accept": "audio/mpeg",
|
||||
"Content-Type": "application/json",
|
||||
"xi-api-key": api_key,
|
||||
}
|
||||
data = {
|
||||
"text": text,
|
||||
"model_id": self.config.get("elevenlabs", {}).get("model"),
|
||||
"voice_settings": {"stability": 0.5, "similarity_boost": 0.5},
|
||||
}
|
||||
|
||||
response = await client.post(url, json=data, headers=headers, timeout=300)
|
||||
|
||||
if response.status_code == 200:
|
||||
bytes_io = io.BytesIO()
|
||||
for chunk in response.iter_bytes(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
bytes_io.write(chunk)
|
||||
|
||||
# Put the audio data in the queue for playback
|
||||
return bytes_io.getvalue()
|
||||
else:
|
||||
log.error(f"Error generating audio: {response.text}")
|
||||
|
||||
async def _list_voices_elevenlabs(self) -> dict[str, str]:
|
||||
url_voices = "https://api.elevenlabs.io/v1/voices"
|
||||
|
||||
voices = []
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"xi-api-key": self.token,
|
||||
}
|
||||
response = await client.get(
|
||||
url_voices, headers=headers, params={"per_page": 1000}
|
||||
)
|
||||
speakers = response.json()["voices"]
|
||||
voices.extend(
|
||||
[
|
||||
Voice(value=speaker["voice_id"], label=speaker["name"])
|
||||
for speaker in speakers
|
||||
]
|
||||
)
|
||||
|
||||
# sort by name
|
||||
voices.sort(key=lambda x: x.label)
|
||||
|
||||
return voices
|
||||
|
||||
# OPENAI
|
||||
|
||||
async def _generate_openai(self, text: str, chunk_size: int = 1024):
|
||||
|
||||
client = AsyncOpenAI(api_key=self.openai_api_key)
|
||||
|
||||
model = self.actions["openai"].config["model"].value
|
||||
|
||||
response = await client.audio.speech.create(
|
||||
model=model, voice=self.default_voice_id, input=text
|
||||
)
|
||||
|
||||
bytes_io = io.BytesIO()
|
||||
for chunk in response.iter_bytes(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
bytes_io.write(chunk)
|
||||
|
||||
# Put the audio data in the queue for playback
|
||||
return bytes_io.getvalue()
|
||||
|
||||
async def _list_voices_openai(self) -> dict[str, str]:
|
||||
return [
|
||||
Voice(value="alloy", label="Alloy"),
|
||||
Voice(value="echo", label="Echo"),
|
||||
Voice(value="fable", label="Fable"),
|
||||
Voice(value="onyx", label="Onyx"),
|
||||
Voice(value="nova", label="Nova"),
|
||||
Voice(value="shimmer", label="Shimmer"),
|
||||
]
|
||||
452
src/talemate/agents/visual/__init__.py
Normal file
@@ -0,0 +1,452 @@
|
||||
import asyncio
|
||||
import traceback
|
||||
|
||||
import structlog
|
||||
|
||||
import talemate.agents.visual.automatic1111
|
||||
import talemate.agents.visual.comfyui
|
||||
import talemate.agents.visual.openai_image
|
||||
from talemate.agents.base import (
|
||||
Agent,
|
||||
AgentAction,
|
||||
AgentActionConditional,
|
||||
AgentActionConfig,
|
||||
AgentDetail,
|
||||
set_processing,
|
||||
)
|
||||
from talemate.agents.registry import register
|
||||
from talemate.client.base import ClientBase
|
||||
from talemate.config import load_config
|
||||
from talemate.emit import emit
|
||||
from talemate.emit.signals import handlers as signal_handlers
|
||||
from talemate.prompts.base import Prompt
|
||||
|
||||
from .commands import * # noqa
|
||||
from .context import VIS_TYPES, VisualContext, visual_context
|
||||
from .handlers import HANDLERS
|
||||
from .schema import RESOLUTION_MAP, RenderSettings
|
||||
from .style import MAJOR_STYLES, STYLE_MAP, Style, combine_styles
|
||||
from .websocket_handler import VisualWebsocketHandler
|
||||
|
||||
__all__ = [
|
||||
"VisualAgent",
|
||||
]
|
||||
|
||||
BACKENDS = [
|
||||
{"value": mixin_backend, "label": mixin["label"]}
|
||||
for mixin_backend, mixin in HANDLERS.items()
|
||||
]
|
||||
|
||||
log = structlog.get_logger("talemate.agents.visual")
|
||||
|
||||
|
||||
class VisualBase(Agent):
|
||||
"""
|
||||
The visual agent
|
||||
"""
|
||||
|
||||
agent_type = "visual"
|
||||
verbose_name = "Visualizer"
|
||||
essential = False
|
||||
websocket_handler = VisualWebsocketHandler
|
||||
|
||||
ACTIONS = {}
|
||||
|
||||
def __init__(self, client: ClientBase, *kwargs):
|
||||
self.client = client
|
||||
self.is_enabled = False
|
||||
self.backend_ready = False
|
||||
self.initialized = False
|
||||
self.config = load_config()
|
||||
self.actions = {
|
||||
"_config": AgentAction(
|
||||
enabled=True,
|
||||
label="Configure",
|
||||
description="Visual agent configuration",
|
||||
config={
|
||||
"backend": AgentActionConfig(
|
||||
type="text",
|
||||
choices=BACKENDS,
|
||||
value="automatic1111",
|
||||
label="Backend",
|
||||
description="The backend to use for visual processing",
|
||||
),
|
||||
"default_style": AgentActionConfig(
|
||||
type="text",
|
||||
value="ink_illustration",
|
||||
choices=MAJOR_STYLES,
|
||||
label="Default Style",
|
||||
description="The default style to use for visual processing",
|
||||
),
|
||||
},
|
||||
),
|
||||
"automatic_generation": AgentAction(
|
||||
enabled=False,
|
||||
label="Automatic Generation",
|
||||
description="Allow automatic generation of visual content",
|
||||
),
|
||||
"process_in_background": AgentAction(
|
||||
enabled=True,
|
||||
label="Process in Background",
|
||||
description="Process renders in the background",
|
||||
),
|
||||
}
|
||||
|
||||
for action_name, action in self.ACTIONS.items():
|
||||
self.actions[action_name] = action
|
||||
|
||||
signal_handlers["config_saved"].connect(self.on_config_saved)
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return self.is_enabled
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
return True
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return False
|
||||
|
||||
@property
|
||||
def backend(self):
|
||||
return self.actions["_config"].config["backend"].value
|
||||
|
||||
@property
|
||||
def backend_name(self):
|
||||
key = self.actions["_config"].config["backend"].value
|
||||
|
||||
for backend in BACKENDS:
|
||||
if backend["value"] == key:
|
||||
return backend["label"]
|
||||
|
||||
@property
|
||||
def default_style(self):
|
||||
return STYLE_MAP.get(
|
||||
self.actions["_config"].config["default_style"].value, Style()
|
||||
)
|
||||
|
||||
@property
|
||||
def ready(self):
|
||||
return self.backend_ready
|
||||
|
||||
@property
|
||||
def api_url(self):
|
||||
try:
|
||||
return self.actions[self.backend].config["api_url"].value
|
||||
except KeyError:
|
||||
return None
|
||||
|
||||
@property
|
||||
def agent_details(self):
|
||||
details = {
|
||||
"backend": AgentDetail(
|
||||
icon="mdi-server-outline",
|
||||
value=self.backend_name,
|
||||
description="The backend to use for visual processing",
|
||||
).model_dump(),
|
||||
"client": AgentDetail(
|
||||
icon="mdi-network-outline",
|
||||
value=self.client.name if self.client else None,
|
||||
description="The client to use for prompt generation",
|
||||
).model_dump(),
|
||||
}
|
||||
|
||||
if not self.ready and self.enabled:
|
||||
details["status"] = AgentDetail(
|
||||
icon="mdi-alert",
|
||||
value=f"{self.backend_name} not ready",
|
||||
color="error",
|
||||
description=self.ready_check_error
|
||||
or f"{self.backend_name} is not ready for processing",
|
||||
).model_dump()
|
||||
|
||||
return details
|
||||
|
||||
@property
|
||||
def process_in_background(self):
|
||||
return self.actions["process_in_background"].enabled
|
||||
|
||||
@property
|
||||
def allow_automatic_generation(self):
|
||||
return self.actions["automatic_generation"].enabled
|
||||
|
||||
def on_config_saved(self, event):
|
||||
config = event.data
|
||||
self.config = config
|
||||
asyncio.create_task(self.emit_status())
|
||||
|
||||
async def on_ready_check_success(self):
|
||||
prev_ready = self.backend_ready
|
||||
self.backend_ready = True
|
||||
if not prev_ready:
|
||||
await self.emit_status()
|
||||
|
||||
async def on_ready_check_failure(self, error):
|
||||
prev_ready = self.backend_ready
|
||||
self.backend_ready = False
|
||||
self.ready_check_error = str(error)
|
||||
if prev_ready:
|
||||
await self.emit_status()
|
||||
|
||||
async def ready_check(self):
|
||||
if not self.enabled:
|
||||
return
|
||||
backend = self.backend
|
||||
fn = getattr(self, f"{backend.lower()}_ready", None)
|
||||
task = asyncio.create_task(fn())
|
||||
await super().ready_check(task)
|
||||
|
||||
async def apply_config(self, *args, **kwargs):
|
||||
|
||||
try:
|
||||
backend = kwargs["actions"]["_config"]["config"]["backend"]["value"]
|
||||
except KeyError:
|
||||
backend = self.backend
|
||||
|
||||
backend_changed = backend != self.backend
|
||||
|
||||
if backend_changed:
|
||||
self.backend_ready = False
|
||||
|
||||
log.info(
|
||||
"apply_config",
|
||||
backend=backend,
|
||||
backend_changed=backend_changed,
|
||||
old_backend=self.backend,
|
||||
)
|
||||
|
||||
await super().apply_config(*args, **kwargs)
|
||||
backend_fn = getattr(self, f"{self.backend.lower()}_apply_config", None)
|
||||
if backend_fn:
|
||||
task = asyncio.create_task(
|
||||
backend_fn(backend_changed=backend_changed, *args, **kwargs)
|
||||
)
|
||||
await self.set_background_processing(task)
|
||||
|
||||
if not self.backend_ready:
|
||||
await self.ready_check()
|
||||
|
||||
self.initialized = True
|
||||
|
||||
def resolution_from_format(self, format: str, model_type: str = "sdxl"):
|
||||
if model_type not in RESOLUTION_MAP:
|
||||
raise ValueError(f"Model type {model_type} not found in resolution map")
|
||||
return RESOLUTION_MAP[model_type].get(
|
||||
format, RESOLUTION_MAP[model_type]["portrait"]
|
||||
)
|
||||
|
||||
def prepare_prompt(self, prompt: str, styles: list[Style] = None) -> Style:
|
||||
|
||||
prompt_style = Style()
|
||||
prompt_style.load(prompt)
|
||||
|
||||
if styles:
|
||||
prompt_style.prepend(*styles)
|
||||
|
||||
return prompt_style
|
||||
|
||||
def vis_type_styles(self, vis_type: str):
|
||||
if vis_type == VIS_TYPES.CHARACTER:
|
||||
portrait_style = STYLE_MAP["character_portrait"].copy()
|
||||
return portrait_style
|
||||
elif vis_type == VIS_TYPES.ENVIRONMENT:
|
||||
environment_style = STYLE_MAP["environment"].copy()
|
||||
return environment_style
|
||||
return Style()
|
||||
|
||||
async def apply_image(self, image: str):
|
||||
context = visual_context.get()
|
||||
|
||||
log.debug("apply_image", image=image[:100], context=context)
|
||||
|
||||
if context.vis_type == VIS_TYPES.CHARACTER:
|
||||
await self.apply_image_character(image, context.character_name)
|
||||
|
||||
async def apply_image_character(self, image: str, character_name: str):
|
||||
character = self.scene.get_character(character_name)
|
||||
|
||||
if not character:
|
||||
log.error("character not found", character_name=character_name)
|
||||
return
|
||||
|
||||
if character.cover_image:
|
||||
log.info("character cover image already set", character_name=character_name)
|
||||
return
|
||||
|
||||
asset = self.scene.assets.add_asset_from_image_data(
|
||||
f"data:image/png;base64,{image}"
|
||||
)
|
||||
character.cover_image = asset.id
|
||||
self.scene.assets.cover_image = asset.id
|
||||
self.scene.emit_status()
|
||||
|
||||
async def emit_image(self, image: str):
|
||||
context = visual_context.get()
|
||||
await self.apply_image(image)
|
||||
emit(
|
||||
"image_generated",
|
||||
websocket_passthrough=True,
|
||||
data={
|
||||
"base64": image,
|
||||
"context": context.model_dump() if context else None,
|
||||
},
|
||||
)
|
||||
|
||||
@set_processing
|
||||
async def generate(
|
||||
self, format: str = "portrait", prompt: str = None, automatic: bool = False
|
||||
):
|
||||
|
||||
context = visual_context.get()
|
||||
|
||||
if not self.enabled:
|
||||
log.warning("generate", skipped="Visual agent not enabled")
|
||||
return
|
||||
|
||||
if automatic and not self.allow_automatic_generation:
|
||||
log.warning(
|
||||
"generate",
|
||||
skipped="Automatic generation disabled",
|
||||
prompt=prompt,
|
||||
format=format,
|
||||
context=context,
|
||||
)
|
||||
return
|
||||
|
||||
if not context and not prompt:
|
||||
log.error("generate", error="No context or prompt provided")
|
||||
return
|
||||
|
||||
# Handle prompt generation based on context
|
||||
|
||||
if not prompt and context.prompt:
|
||||
prompt = context.prompt
|
||||
|
||||
if context.vis_type == VIS_TYPES.ENVIRONMENT and not prompt:
|
||||
prompt = await self.generate_environment_prompt(
|
||||
instructions=context.instructions
|
||||
)
|
||||
elif context.vis_type == VIS_TYPES.CHARACTER and not prompt:
|
||||
prompt = await self.generate_character_prompt(
|
||||
context.character_name, instructions=context.instructions
|
||||
)
|
||||
else:
|
||||
prompt = prompt or context.prompt
|
||||
|
||||
initial_prompt = prompt
|
||||
|
||||
# Augment the prompt with styles based on context
|
||||
|
||||
thematic_style = self.default_style
|
||||
vis_type_styles = self.vis_type_styles(context.vis_type)
|
||||
prompt = self.prepare_prompt(prompt, [vis_type_styles, thematic_style])
|
||||
|
||||
if not prompt:
|
||||
log.error(
|
||||
"generate", error="No prompt provided and no context to generate from"
|
||||
)
|
||||
return
|
||||
|
||||
context.prompt = initial_prompt
|
||||
context.prepared_prompt = str(prompt)
|
||||
|
||||
# Handle format (can either come from context or be passed in)
|
||||
|
||||
if not format and context.format:
|
||||
format = context.format
|
||||
elif not format:
|
||||
format = "portrait"
|
||||
|
||||
context.format = format
|
||||
|
||||
# Call the backend specific generate function
|
||||
|
||||
backend = self.backend
|
||||
fn = f"{backend.lower()}_generate"
|
||||
|
||||
log.info(
|
||||
"generate", backend=backend, prompt=prompt, format=format, context=context
|
||||
)
|
||||
|
||||
if not hasattr(self, fn):
|
||||
log.error("generate", error=f"Backend {backend} does not support generate")
|
||||
|
||||
# add the function call to the asyncio task queue
|
||||
|
||||
if self.process_in_background:
|
||||
task = asyncio.create_task(getattr(self, fn)(prompt=prompt, format=format))
|
||||
await self.set_background_processing(task)
|
||||
else:
|
||||
await getattr(self, fn)(prompt=prompt, format=format)
|
||||
|
||||
@set_processing
|
||||
async def generate_environment_prompt(self, instructions: str = None):
|
||||
|
||||
response = await Prompt.request(
|
||||
"visual.generate-environment-prompt",
|
||||
self.client,
|
||||
"visualize",
|
||||
{
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
},
|
||||
)
|
||||
|
||||
return response.strip()
|
||||
|
||||
@set_processing
|
||||
async def generate_character_prompt(
|
||||
self, character_name: str, instructions: str = None
|
||||
):
|
||||
|
||||
character = self.scene.get_character(character_name)
|
||||
|
||||
response = await Prompt.request(
|
||||
"visual.generate-character-prompt",
|
||||
self.client,
|
||||
"visualize",
|
||||
{
|
||||
"scene": self.scene,
|
||||
"character_name": character_name,
|
||||
"character": character,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"instructions": instructions or "",
|
||||
},
|
||||
)
|
||||
|
||||
return response.strip()
|
||||
|
||||
async def generate_environment_background(self, instructions: str = None):
|
||||
with VisualContext(vis_type=VIS_TYPES.ENVIRONMENT, instructions=instructions):
|
||||
await self.generate(format="landscape")
|
||||
|
||||
async def generate_character_portrait(
|
||||
self,
|
||||
character_name: str,
|
||||
instructions: str = None,
|
||||
):
|
||||
with VisualContext(
|
||||
vis_type=VIS_TYPES.CHARACTER,
|
||||
character_name=character_name,
|
||||
instructions=instructions,
|
||||
):
|
||||
await self.generate(format="portrait")
|
||||
|
||||
|
||||
# apply mixins to the agent (from HANDLERS dict[str, cls])
|
||||
|
||||
for mixin_backend, mixin in HANDLERS.items():
|
||||
mixin_cls = mixin["cls"]
|
||||
VisualBase = type("VisualAgent", (mixin_cls, VisualBase), {})
|
||||
|
||||
extend_actions = getattr(mixin_cls, "EXTEND_ACTIONS", {})
|
||||
|
||||
for action_name, action in extend_actions.items():
|
||||
VisualBase.ACTIONS[action_name] = action
|
||||
|
||||
|
||||
@register()
|
||||
class VisualAgent(VisualBase):
|
||||
pass
|
||||
117
src/talemate/agents/visual/automatic1111.py
Normal file
@@ -0,0 +1,117 @@
|
||||
import base64
|
||||
import io
|
||||
|
||||
import httpx
|
||||
import structlog
|
||||
from PIL import Image
|
||||
|
||||
from talemate.agents.base import (
|
||||
Agent,
|
||||
AgentAction,
|
||||
AgentActionConditional,
|
||||
AgentActionConfig,
|
||||
AgentDetail,
|
||||
set_processing,
|
||||
)
|
||||
|
||||
from .handlers import register
|
||||
from .schema import RenderSettings, Resolution
|
||||
from .style import STYLE_MAP, Style
|
||||
|
||||
log = structlog.get_logger("talemate.agents.visual.automatic1111")
|
||||
|
||||
|
||||
@register(backend_name="automatic1111", label="AUTOMATIC1111")
|
||||
class Automatic1111Mixin:
|
||||
|
||||
automatic1111_default_render_settings = RenderSettings()
|
||||
|
||||
EXTEND_ACTIONS = {
|
||||
"automatic1111": AgentAction(
|
||||
enabled=True,
|
||||
condition=AgentActionConditional(
|
||||
attribute="_config.config.backend", value="automatic1111"
|
||||
),
|
||||
label="Automatic1111 Settings",
|
||||
description="Setting overrides for the automatic1111 backend",
|
||||
config={
|
||||
"api_url": AgentActionConfig(
|
||||
type="text",
|
||||
value="http://localhost:7860",
|
||||
label="API URL",
|
||||
description="The URL of the backend API",
|
||||
),
|
||||
"steps": AgentActionConfig(
|
||||
type="number",
|
||||
value=40,
|
||||
label="Steps",
|
||||
min=5,
|
||||
max=150,
|
||||
step=1,
|
||||
description="number of render steps",
|
||||
),
|
||||
"model_type": AgentActionConfig(
|
||||
type="text",
|
||||
value="sdxl",
|
||||
choices=[
|
||||
{"value": "sdxl", "label": "SDXL"},
|
||||
{"value": "sd15", "label": "SD1.5"},
|
||||
],
|
||||
label="Model Type",
|
||||
description="Right now just differentiates between sdxl and sd15 - affect generation resolution",
|
||||
),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@property
|
||||
def automatic1111_render_settings(self):
|
||||
if self.actions["automatic1111"].enabled:
|
||||
return RenderSettings(
|
||||
steps=self.actions["automatic1111"].config["steps"].value,
|
||||
type_model=self.actions["automatic1111"].config["model_type"].value,
|
||||
)
|
||||
else:
|
||||
return self.automatic1111_default_render_settings
|
||||
|
||||
async def automatic1111_generate(self, prompt: Style, format: str):
|
||||
url = self.api_url
|
||||
resolution = self.resolution_from_format(
|
||||
format, self.automatic1111_render_settings.type_model
|
||||
)
|
||||
render_settings = self.automatic1111_render_settings
|
||||
payload = {
|
||||
"prompt": prompt.positive_prompt,
|
||||
"negative_prompt": prompt.negative_prompt,
|
||||
"steps": render_settings.steps,
|
||||
"width": resolution.width,
|
||||
"height": resolution.height,
|
||||
}
|
||||
|
||||
log.info("automatic1111_generate", payload=payload, url=url)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
url=f"{url}/sdapi/v1/txt2img", json=payload, timeout=90
|
||||
)
|
||||
|
||||
r = response.json()
|
||||
|
||||
# image = Image.open(io.BytesIO(base64.b64decode(r['images'][0])))
|
||||
# image.save('a1111-test.png')
|
||||
|
||||
#'log.info("automatic1111_generate", saved_to="a1111-test.png")
|
||||
|
||||
for image in r["images"]:
|
||||
await self.emit_image(image)
|
||||
|
||||
async def automatic1111_ready(self) -> bool:
|
||||
"""
|
||||
Will send a GET to /sdapi/v1/memory and on 200 will return True
|
||||
"""
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
url=f"{self.api_url}/sdapi/v1/memory", timeout=2
|
||||
)
|
||||
return response.status_code == 200
|
||||
324
src/talemate/agents/visual/comfyui.py
Normal file
@@ -0,0 +1,324 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import urllib.parse
|
||||
|
||||
import httpx
|
||||
import pydantic
|
||||
import structlog
|
||||
from PIL import Image
|
||||
|
||||
from talemate.agents.base import AgentAction, AgentActionConditional, AgentActionConfig
|
||||
|
||||
from .handlers import register
|
||||
from .schema import RenderSettings, Resolution
|
||||
from .style import STYLE_MAP, Style
|
||||
|
||||
log = structlog.get_logger("talemate.agents.visual.comfyui")
|
||||
|
||||
|
||||
class Workflow(pydantic.BaseModel):
|
||||
nodes: dict
|
||||
|
||||
def set_resolution(self, resolution: Resolution):
|
||||
|
||||
# will collect all latent image nodes
|
||||
# if there is multiple will look for the one with the
|
||||
# title "Talemate Resolution"
|
||||
|
||||
# if there is no latent image node with the title "Talemate Resolution"
|
||||
# the first latent image node will be used
|
||||
|
||||
# resolution will be updated on the selected node
|
||||
|
||||
# if no latent image node is found a warning will be logged
|
||||
|
||||
latent_image_node = None
|
||||
|
||||
for node_id, node in self.nodes.items():
|
||||
if node["class_type"] == "EmptyLatentImage":
|
||||
if not latent_image_node:
|
||||
latent_image_node = node
|
||||
elif node["_meta"]["title"] == "Talemate Resolution":
|
||||
latent_image_node = node
|
||||
break
|
||||
|
||||
if not latent_image_node:
|
||||
log.warning("set_resolution", error="No latent image node found")
|
||||
return
|
||||
|
||||
latent_image_node["inputs"]["width"] = resolution.width
|
||||
latent_image_node["inputs"]["height"] = resolution.height
|
||||
|
||||
def set_prompt(self, prompt: str, negative_prompt: str = None):
|
||||
|
||||
# will collect all CLIPTextEncode nodes
|
||||
|
||||
# if there is multiple will look for the one with the
|
||||
# title "Talemate Positive Prompt" and "Talemate Negative Prompt"
|
||||
#
|
||||
# if there is no CLIPTextEncode node with the title "Talemate Positive Prompt"
|
||||
# the first CLIPTextEncode node will be used
|
||||
#
|
||||
# if there is no CLIPTextEncode node with the title "Talemate Negative Prompt"
|
||||
# the second CLIPTextEncode node will be used
|
||||
#
|
||||
# prompt will be updated on the selected node
|
||||
|
||||
# if no CLIPTextEncode node is found an exception will be raised for
|
||||
# the positive prompt
|
||||
|
||||
# if no CLIPTextEncode node is found an exception will be raised for
|
||||
# the negative prompt if it is not None
|
||||
|
||||
positive_prompt_node = None
|
||||
negative_prompt_node = None
|
||||
|
||||
for node_id, node in self.nodes.items():
|
||||
|
||||
if node["class_type"] == "CLIPTextEncode":
|
||||
if not positive_prompt_node:
|
||||
positive_prompt_node = node
|
||||
elif node["_meta"]["title"] == "Talemate Positive Prompt":
|
||||
positive_prompt_node = node
|
||||
elif not negative_prompt_node:
|
||||
negative_prompt_node = node
|
||||
elif node["_meta"]["title"] == "Talemate Negative Prompt":
|
||||
negative_prompt_node = node
|
||||
|
||||
if not positive_prompt_node:
|
||||
raise ValueError("No positive prompt node found")
|
||||
|
||||
positive_prompt_node["inputs"]["text"] = prompt
|
||||
|
||||
if negative_prompt and not negative_prompt_node:
|
||||
raise ValueError("No negative prompt node found")
|
||||
|
||||
if negative_prompt:
|
||||
negative_prompt_node["inputs"]["text"] = negative_prompt
|
||||
|
||||
def set_checkpoint(self, checkpoint: str):
|
||||
|
||||
# will collect all CheckpointLoaderSimple nodes
|
||||
# if there is multiple will look for the one with the
|
||||
# title "Talemate Load Checkpoint"
|
||||
|
||||
# if there is no CheckpointLoaderSimple node with the title "Talemate Load Checkpoint"
|
||||
# the first CheckpointLoaderSimple node will be used
|
||||
|
||||
# checkpoint will be updated on the selected node
|
||||
|
||||
# if no CheckpointLoaderSimple node is found a warning will be logged
|
||||
|
||||
checkpoint_node = None
|
||||
|
||||
for node_id, node in self.nodes.items():
|
||||
if node["class_type"] == "CheckpointLoaderSimple":
|
||||
if not checkpoint_node:
|
||||
checkpoint_node = node
|
||||
elif node["_meta"]["title"] == "Talemate Load Checkpoint":
|
||||
checkpoint_node = node
|
||||
break
|
||||
|
||||
if not checkpoint_node:
|
||||
log.warning("set_checkpoint", error="No checkpoint node found")
|
||||
return
|
||||
|
||||
checkpoint_node["inputs"]["ckpt_name"] = checkpoint
|
||||
|
||||
def set_seeds(self):
|
||||
for node in self.nodes.values():
|
||||
for field in node.get("inputs", {}).keys():
|
||||
if field == "noise_seed":
|
||||
node["inputs"]["noise_seed"] = random.randint(0, 999999999999999)
|
||||
|
||||
|
||||
@register(backend_name="comfyui", label="ComfyUI")
|
||||
class ComfyUIMixin:
|
||||
|
||||
comfyui_default_render_settings = RenderSettings()
|
||||
|
||||
EXTEND_ACTIONS = {
|
||||
"comfyui": AgentAction(
|
||||
enabled=True,
|
||||
condition=AgentActionConditional(
|
||||
attribute="_config.config.backend", value="comfyui"
|
||||
),
|
||||
label="ComfyUI Settings",
|
||||
description="Setting overrides for the comfyui backend",
|
||||
config={
|
||||
"api_url": AgentActionConfig(
|
||||
type="text",
|
||||
value="http://localhost:8188",
|
||||
label="API URL",
|
||||
description="The URL of the backend API",
|
||||
),
|
||||
"workflow": AgentActionConfig(
|
||||
type="text",
|
||||
value="default-sdxl.json",
|
||||
label="Workflow",
|
||||
description="The workflow to use for comfyui (workflow file name inside ./templates/comfyui-workflows)",
|
||||
),
|
||||
"checkpoint": AgentActionConfig(
|
||||
type="text",
|
||||
value="default",
|
||||
label="Checkpoint",
|
||||
choices=[],
|
||||
description="The main checkpoint to use.",
|
||||
),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@property
|
||||
def comfyui_workflow_filename(self):
|
||||
base_name = self.actions["comfyui"].config["workflow"].value
|
||||
|
||||
# make absolute path
|
||||
abs_path = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"templates",
|
||||
"comfyui-workflows",
|
||||
base_name,
|
||||
)
|
||||
|
||||
return abs_path
|
||||
|
||||
@property
|
||||
def comfyui_workflow_is_sdxl(self) -> bool:
|
||||
"""
|
||||
Returns true if `sdxl` is in worhflow file name (case insensitive)
|
||||
"""
|
||||
|
||||
return "sdxl" in self.comfyui_workflow_filename.lower()
|
||||
|
||||
@property
|
||||
def comfyui_workflow(self) -> Workflow:
|
||||
workflow = self.comfyui_workflow_filename
|
||||
if not workflow:
|
||||
raise ValueError("No comfyui workflow file specified")
|
||||
|
||||
with open(workflow, "r") as f:
|
||||
return Workflow(nodes=json.load(f))
|
||||
|
||||
@property
|
||||
async def comfyui_object_info(self):
|
||||
if hasattr(self, "_comfyui_object_info"):
|
||||
return self._comfyui_object_info
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url=f"{self.api_url}/object_info")
|
||||
self._comfyui_object_info = response.json()
|
||||
|
||||
return self._comfyui_object_info
|
||||
|
||||
@property
|
||||
async def comfyui_checkpoints(self):
|
||||
loader_node = (await self.comfyui_object_info)["CheckpointLoaderSimple"]
|
||||
_checkpoints = loader_node["input"]["required"]["ckpt_name"][0]
|
||||
return [
|
||||
{"label": checkpoint, "value": checkpoint} for checkpoint in _checkpoints
|
||||
]
|
||||
|
||||
async def comfyui_get_image(self, filename: str, subfolder: str, folder_type: str):
|
||||
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
||||
url_values = urllib.parse.urlencode(data)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url=f"{self.api_url}/view?{url_values}")
|
||||
return response.content
|
||||
|
||||
async def comfyui_get_history(self, prompt_id: str):
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url=f"{self.api_url}/history/{prompt_id}")
|
||||
return response.json()
|
||||
|
||||
async def comfyui_get_images(self, prompt_id: str, max_wait: int = 60.0):
|
||||
output_images = {}
|
||||
history = {}
|
||||
|
||||
start = time.time()
|
||||
|
||||
while not history:
|
||||
log.info(
|
||||
"comfyui_get_images", waiting_for_history=True, prompt_id=prompt_id
|
||||
)
|
||||
history = await self.comfyui_get_history(prompt_id)
|
||||
await asyncio.sleep(1.0)
|
||||
if time.time() - start > max_wait:
|
||||
raise TimeoutError("Max wait time exceeded")
|
||||
|
||||
for node_id, node_output in history[prompt_id]["outputs"].items():
|
||||
if "images" in node_output:
|
||||
images_output = []
|
||||
for image in node_output["images"]:
|
||||
image_data = await self.comfyui_get_image(
|
||||
image["filename"], image["subfolder"], image["type"]
|
||||
)
|
||||
images_output.append(image_data)
|
||||
output_images[node_id] = images_output
|
||||
|
||||
return output_images
|
||||
|
||||
async def comfyui_generate(self, prompt: Style, format: str):
|
||||
url = self.api_url
|
||||
workflow = self.comfyui_workflow
|
||||
is_sdxl = self.comfyui_workflow_is_sdxl
|
||||
|
||||
resolution = self.resolution_from_format(format, "sdxl" if is_sdxl else "sd15")
|
||||
|
||||
workflow.set_resolution(resolution)
|
||||
workflow.set_prompt(prompt.positive_prompt, prompt.negative_prompt)
|
||||
workflow.set_seeds()
|
||||
workflow.set_checkpoint(self.actions["comfyui"].config["checkpoint"].value)
|
||||
|
||||
payload = {"prompt": workflow.model_dump().get("nodes")}
|
||||
|
||||
log.info("comfyui_generate", payload=payload, url=url)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(url=f"{url}/prompt", json=payload, timeout=90)
|
||||
|
||||
log.info("comfyui_generate", response=response.text)
|
||||
|
||||
r = response.json()
|
||||
|
||||
prompt_id = r["prompt_id"]
|
||||
|
||||
images = await self.comfyui_get_images(prompt_id)
|
||||
for node_id, node_images in images.items():
|
||||
for i, image in enumerate(node_images):
|
||||
await self.emit_image(base64.b64encode(image).decode("utf-8"))
|
||||
# image = Image.open(io.BytesIO(image))
|
||||
# image.save(f'comfyui-test.png')
|
||||
|
||||
async def comfyui_apply_config(
|
||||
self, backend_changed: bool = False, *args, **kwargs
|
||||
):
|
||||
log.debug(
|
||||
"comfyui_apply_config",
|
||||
backend_changed=backend_changed,
|
||||
enabled=self.enabled,
|
||||
)
|
||||
if (not self.initialized or backend_changed) and self.enabled:
|
||||
checkpoints = await self.comfyui_checkpoints
|
||||
selected_checkpoint = self.actions["comfyui"].config["checkpoint"].value
|
||||
self.actions["comfyui"].config["checkpoint"].choices = checkpoints
|
||||
self.actions["comfyui"].config["checkpoint"].value = selected_checkpoint
|
||||
|
||||
async def comfyui_ready(self) -> bool:
|
||||
"""
|
||||
Will send a GET to /system_stats and on 200 will return True
|
||||
"""
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url=f"{self.api_url}/system_stats", timeout=2)
|
||||
return response.status_code == 200
|
||||
68
src/talemate/agents/visual/commands.py
Normal file
@@ -0,0 +1,68 @@
|
||||
from talemate.agents.visual.context import VIS_TYPES, VisualContext
|
||||
from talemate.commands.base import TalemateCommand
|
||||
from talemate.commands.manager import register
|
||||
from talemate.instance import get_agent
|
||||
|
||||
__all__ = [
|
||||
"CmdVisualizeTestGenerate",
|
||||
]
|
||||
|
||||
|
||||
@register
|
||||
class CmdVisualizeTestGenerate(TalemateCommand):
|
||||
"""
|
||||
Generates a visual test
|
||||
"""
|
||||
|
||||
name = "visual_test_generate"
|
||||
description = "Will generate a visual test"
|
||||
aliases = ["vis_test", "vtg"]
|
||||
|
||||
label = "Visualize test"
|
||||
|
||||
async def run(self):
|
||||
visual = get_agent("visual")
|
||||
prompt = self.args[0]
|
||||
with VisualContext(vis_type=VIS_TYPES.UNSPECIFIED):
|
||||
await visual.generate(prompt)
|
||||
return True
|
||||
|
||||
|
||||
@register
|
||||
class CmdVisualizeEnvironment(TalemateCommand):
|
||||
"""
|
||||
Shows the environment
|
||||
"""
|
||||
|
||||
name = "visual_environment"
|
||||
description = "Will show the environment"
|
||||
aliases = ["vis_env"]
|
||||
|
||||
label = "Visualize environment"
|
||||
|
||||
async def run(self):
|
||||
visual = get_agent("visual")
|
||||
await visual.generate_environment_background(
|
||||
instructions=self.args[0] if len(self.args) > 0 else None
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@register
|
||||
class CmdVisualizeCharacter(TalemateCommand):
|
||||
"""
|
||||
Shows a character
|
||||
"""
|
||||
|
||||
name = "visual_character"
|
||||
description = "Will show a character"
|
||||
aliases = ["vis_char"]
|
||||
|
||||
label = "Visualize character"
|
||||
|
||||
async def run(self):
|
||||
visual = get_agent("visual")
|
||||
character_name = self.args[0]
|
||||
instructions = self.args[1] if len(self.args) > 1 else None
|
||||
await visual.generate_character_portrait(character_name, instructions)
|
||||
return True
|
||||
55
src/talemate/agents/visual/context.py
Normal file
@@ -0,0 +1,55 @@
|
||||
import contextvars
|
||||
import enum
|
||||
from typing import Union
|
||||
|
||||
import pydantic
|
||||
|
||||
__all__ = [
|
||||
"VIS_TYPES",
|
||||
"visual_context",
|
||||
"VisualContext",
|
||||
]
|
||||
|
||||
|
||||
class VIS_TYPES(str, enum.Enum):
|
||||
UNSPECIFIED = "UNSPECIFIED"
|
||||
ENVIRONMENT = "ENVIRONMENT"
|
||||
CHARACTER = "CHARACTER"
|
||||
ITEM = "ITEM"
|
||||
|
||||
|
||||
visual_context = contextvars.ContextVar("visual_context", default=None)
|
||||
|
||||
|
||||
class VisualContextState(pydantic.BaseModel):
|
||||
character_name: Union[str, None] = None
|
||||
instructions: Union[str, None] = None
|
||||
vis_type: VIS_TYPES = VIS_TYPES.ENVIRONMENT
|
||||
prompt: Union[str, None] = None
|
||||
prepared_prompt: Union[str, None] = None
|
||||
format: Union[str, None] = None
|
||||
|
||||
|
||||
class VisualContext:
|
||||
def __init__(
|
||||
self,
|
||||
character_name: Union[str, None] = None,
|
||||
instructions: Union[str, None] = None,
|
||||
vis_type: VIS_TYPES = VIS_TYPES.ENVIRONMENT,
|
||||
prompt: Union[str, None] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.state = VisualContextState(
|
||||
character_name=character_name,
|
||||
instructions=instructions,
|
||||
vis_type=vis_type,
|
||||
prompt=prompt,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def __enter__(self):
|
||||
self.token = visual_context.set(self.state)
|
||||
|
||||
def __exit__(self, *args, **kwargs):
|
||||
visual_context.reset(self.token)
|
||||
return False
|
||||
17
src/talemate/agents/visual/handlers.py
Normal file
@@ -0,0 +1,17 @@
|
||||
__all__ = [
|
||||
"HANDLERS",
|
||||
"register",
|
||||
]
|
||||
|
||||
HANDLERS = {}
|
||||
|
||||
|
||||
class register:
|
||||
|
||||
def __init__(self, backend_name: str, label: str):
|
||||
self.backend_name = backend_name
|
||||
self.label = label
|
||||
|
||||
def __call__(self, mixin_cls):
|
||||
HANDLERS[self.backend_name] = {"label": self.label, "cls": mixin_cls}
|
||||
return mixin_cls
|
||||
127
src/talemate/agents/visual/openai_image.py
Normal file
@@ -0,0 +1,127 @@
|
||||
import base64
|
||||
import io
|
||||
|
||||
import httpx
|
||||
import structlog
|
||||
from openai import AsyncOpenAI
|
||||
from PIL import Image
|
||||
|
||||
from talemate.agents.base import (
|
||||
Agent,
|
||||
AgentAction,
|
||||
AgentActionConditional,
|
||||
AgentActionConfig,
|
||||
AgentDetail,
|
||||
set_processing,
|
||||
)
|
||||
|
||||
from .handlers import register
|
||||
from .schema import RenderSettings, Resolution
|
||||
from .style import STYLE_MAP, Style
|
||||
|
||||
log = structlog.get_logger("talemate.agents.visual.openai_image")
|
||||
|
||||
|
||||
@register(backend_name="openai_image", label="OpenAI")
|
||||
class OpenAIImageMixin:
|
||||
|
||||
openai_image_default_render_settings = RenderSettings()
|
||||
|
||||
EXTEND_ACTIONS = {
|
||||
"openai_image": AgentAction(
|
||||
enabled=False,
|
||||
condition=AgentActionConditional(
|
||||
attribute="_config.config.backend", value="openai_image"
|
||||
),
|
||||
label="OpenAI Image Generation Advanced Settings",
|
||||
description="Setting overrides for the openai backend",
|
||||
config={
|
||||
"model_type": AgentActionConfig(
|
||||
type="text",
|
||||
value="dall-e-3",
|
||||
choices=[
|
||||
{"value": "dall-e-3", "label": "DALL-E 3"},
|
||||
{"value": "dall-e-2", "label": "DALL-E 2"},
|
||||
],
|
||||
label="Model Type",
|
||||
description="Image generation model",
|
||||
),
|
||||
"quality": AgentActionConfig(
|
||||
type="text",
|
||||
value="standard",
|
||||
choices=[
|
||||
{"value": "standard", "label": "Standard"},
|
||||
{"value": "hd", "label": "HD"},
|
||||
],
|
||||
label="Quality",
|
||||
description="Image generation quality",
|
||||
),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@property
|
||||
def openai_api_key(self):
|
||||
return self.config.get("openai", {}).get("api_key")
|
||||
|
||||
@property
|
||||
def openai_model_type(self):
|
||||
return self.actions["openai_image"].config["model_type"].value
|
||||
|
||||
@property
|
||||
def openai_quality(self):
|
||||
return self.actions["openai_image"].config["quality"].value
|
||||
|
||||
async def openai_image_generate(self, prompt: Style, format: str):
|
||||
"""
|
||||
#
|
||||
from openai import OpenAI
|
||||
client = OpenAI()
|
||||
|
||||
response = client.images.generate(
|
||||
model="dall-e-3",
|
||||
prompt="a white siamese cat",
|
||||
size="1024x1024",
|
||||
quality="standard",
|
||||
n=1,
|
||||
)
|
||||
|
||||
image_url = response.data[0].url
|
||||
"""
|
||||
|
||||
client = AsyncOpenAI(api_key=self.openai_api_key)
|
||||
|
||||
# When using DALL·E 3, images can have a size of 1024x1024, 1024x1792 or 1792x1024 pixels.#
|
||||
|
||||
if format == "portrait":
|
||||
resolution = Resolution(width=1024, height=1792)
|
||||
elif format == "landscape":
|
||||
resolution = Resolution(width=1792, height=1024)
|
||||
else:
|
||||
resolution = Resolution(width=1024, height=1024)
|
||||
|
||||
response = await client.images.generate(
|
||||
model=self.openai_model_type,
|
||||
prompt=prompt.positive_prompt,
|
||||
size=f"{resolution.width}x{resolution.height}",
|
||||
quality=self.openai_quality,
|
||||
n=1,
|
||||
)
|
||||
|
||||
download_url = response.data[0].url
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(download_url, timeout=90)
|
||||
# bytes to base64encoded
|
||||
image = base64.b64encode(response.content).decode("utf-8")
|
||||
await self.emit_image(image)
|
||||
|
||||
async def openai_image_ready(self) -> bool:
|
||||
"""
|
||||
Will send a GET to /sdapi/v1/memory and on 200 will return True
|
||||
"""
|
||||
|
||||
if not self.openai_api_key:
|
||||
raise ValueError("OpenAI API Key not set")
|
||||
|
||||
return True
|
||||
32
src/talemate/agents/visual/schema.py
Normal file
@@ -0,0 +1,32 @@
|
||||
import pydantic
|
||||
|
||||
__all__ = [
|
||||
"RenderSettings",
|
||||
"Resolution",
|
||||
"RESOLUTION_MAP",
|
||||
]
|
||||
|
||||
RESOLUTION_MAP = {}
|
||||
|
||||
|
||||
class RenderSettings(pydantic.BaseModel):
|
||||
type_model: str = "sdxl"
|
||||
steps: int = 40
|
||||
|
||||
|
||||
class Resolution(pydantic.BaseModel):
|
||||
width: int
|
||||
height: int
|
||||
|
||||
|
||||
RESOLUTION_MAP["sdxl"] = {
|
||||
"portrait": Resolution(width=832, height=1216),
|
||||
"landscape": Resolution(width=1216, height=832),
|
||||
"square": Resolution(width=1024, height=1024),
|
||||
}
|
||||
|
||||
RESOLUTION_MAP["sd15"] = {
|
||||
"portrait": Resolution(width=512, height=768),
|
||||
"landscape": Resolution(width=768, height=512),
|
||||
"square": Resolution(width=768, height=768),
|
||||
}
|
||||
112
src/talemate/agents/visual/style.py
Normal file
@@ -0,0 +1,112 @@
|
||||
import pydantic
|
||||
|
||||
__all__ = [
|
||||
"Style",
|
||||
"STYLE_MAP",
|
||||
"THEME_MAP",
|
||||
"MAJOR_STYLES",
|
||||
"combine_styles",
|
||||
]
|
||||
|
||||
STYLE_MAP = {}
|
||||
THEME_MAP = {}
|
||||
MAJOR_STYLES = {}
|
||||
|
||||
|
||||
class Style(pydantic.BaseModel):
|
||||
keywords: list[str] = pydantic.Field(default_factory=list)
|
||||
negative_keywords: list[str] = pydantic.Field(default_factory=list)
|
||||
|
||||
@property
|
||||
def positive_prompt(self):
|
||||
return ", ".join(self.keywords)
|
||||
|
||||
@property
|
||||
def negative_prompt(self):
|
||||
return ", ".join(self.negative_keywords)
|
||||
|
||||
def __str__(self):
|
||||
return f"POSITIVE: {self.positive_prompt}\nNEGATIVE: {self.negative_prompt}"
|
||||
|
||||
def load(self, prompt: str, negative_prompt: str = ""):
|
||||
self.keywords = prompt.split(", ")
|
||||
self.negative_keywords = negative_prompt.split(", ")
|
||||
return self
|
||||
|
||||
def prepend(self, *styles):
|
||||
for style in styles:
|
||||
for idx in range(len(style.keywords) - 1, -1, -1):
|
||||
kw = style.keywords[idx]
|
||||
if kw not in self.keywords:
|
||||
self.keywords.insert(0, kw)
|
||||
|
||||
for idx in range(len(style.negative_keywords) - 1, -1, -1):
|
||||
kw = style.negative_keywords[idx]
|
||||
if kw not in self.negative_keywords:
|
||||
self.negative_keywords.insert(0, kw)
|
||||
|
||||
return self
|
||||
|
||||
def append(self, *styles):
|
||||
for style in styles:
|
||||
for kw in style.keywords:
|
||||
if kw not in self.keywords:
|
||||
self.keywords.append(kw)
|
||||
|
||||
for kw in style.negative_keywords:
|
||||
if kw not in self.negative_keywords:
|
||||
self.negative_keywords.append(kw)
|
||||
|
||||
return self
|
||||
|
||||
def copy(self):
|
||||
return Style(
|
||||
keywords=self.keywords.copy(),
|
||||
negative_keywords=self.negative_keywords.copy(),
|
||||
)
|
||||
|
||||
|
||||
# Almost taken straight from some of the fooocus style presets, credit goes to the original author
|
||||
|
||||
STYLE_MAP["digital_art"] = Style(
|
||||
keywords="digital artwork, masterpiece, best quality, high detail".split(", "),
|
||||
negative_keywords="text, watermark, low quality, blurry, photo".split(", "),
|
||||
)
|
||||
|
||||
STYLE_MAP["concept_art"] = Style(
|
||||
keywords="concept art, conceptual sketch, masterpiece, best quality, high detail".split(
|
||||
", "
|
||||
),
|
||||
negative_keywords="text, watermark, low quality, blurry, photo".split(", "),
|
||||
)
|
||||
|
||||
STYLE_MAP["ink_illustration"] = Style(
|
||||
keywords="ink illustration, painting, masterpiece, best quality".split(", "),
|
||||
negative_keywords="text, watermark, low quality, blurry, photo".split(", "),
|
||||
)
|
||||
|
||||
STYLE_MAP["anime"] = Style(
|
||||
keywords="anime, masterpiece, best quality, illustration".split(", "),
|
||||
negative_keywords="text, watermark, low quality, blurry, photo, 3d".split(", "),
|
||||
)
|
||||
|
||||
STYLE_MAP["character_portrait"] = Style(keywords="solo, looking at viewer".split(", "))
|
||||
|
||||
STYLE_MAP["environment"] = Style(
|
||||
keywords="scenery, environment, background, postcard".split(", "),
|
||||
negative_keywords="character, portrait, looking at viewer, people".split(", "),
|
||||
)
|
||||
|
||||
MAJOR_STYLES = [
|
||||
{"value": "digital_art", "label": "Digital Art"},
|
||||
{"value": "concept_art", "label": "Concept Art"},
|
||||
{"value": "ink_illustration", "label": "Ink Illustration"},
|
||||
{"value": "anime", "label": "Anime"},
|
||||
]
|
||||
|
||||
|
||||
def combine_styles(*styles):
|
||||
keywords = []
|
||||
for style in styles:
|
||||
keywords.extend(style.keywords)
|
||||
return Style(keywords=list(set(keywords)))
|
||||
84
src/talemate/agents/visual/websocket_handler.py
Normal file
@@ -0,0 +1,84 @@
|
||||
from typing import Union
|
||||
|
||||
import pydantic
|
||||
import structlog
|
||||
|
||||
from talemate.instance import get_agent
|
||||
from talemate.server.websocket_plugin import Plugin
|
||||
|
||||
from .context import VisualContext, VisualContextState
|
||||
|
||||
__all__ = [
|
||||
"VisualWebsocketHandler",
|
||||
]
|
||||
|
||||
log = structlog.get_logger("talemate.server.visual")
|
||||
|
||||
|
||||
class SetCoverImagePayload(pydantic.BaseModel):
|
||||
base64: str
|
||||
context: Union[VisualContextState, None] = None
|
||||
|
||||
|
||||
class RegeneratePayload(pydantic.BaseModel):
|
||||
context: Union[VisualContextState, None] = None
|
||||
|
||||
|
||||
class VisualWebsocketHandler(Plugin):
|
||||
router = "visual"
|
||||
|
||||
async def handle_regenerate(self, data: dict):
|
||||
"""
|
||||
Regenerate the image based on the context.
|
||||
"""
|
||||
|
||||
payload = RegeneratePayload(**data)
|
||||
|
||||
context = payload.context
|
||||
|
||||
visual = get_agent("visual")
|
||||
|
||||
with VisualContext(**context.model_dump()):
|
||||
await visual.generate(format="")
|
||||
|
||||
async def handle_cover_image(self, data: dict):
|
||||
"""
|
||||
Sets the cover image for a character and the scene.
|
||||
"""
|
||||
|
||||
payload = SetCoverImagePayload(**data)
|
||||
|
||||
context = payload.context
|
||||
scene = self.scene
|
||||
|
||||
if context and context.character_name:
|
||||
|
||||
character = scene.get_character(context.character_name)
|
||||
|
||||
if not character:
|
||||
log.error("character not found", character_name=context.character_name)
|
||||
return
|
||||
|
||||
asset = scene.assets.add_asset_from_image_data(payload.base64)
|
||||
|
||||
log.info("setting scene cover image", character_name=context.character_name)
|
||||
scene.assets.cover_image = asset.id
|
||||
|
||||
log.info(
|
||||
"setting character cover image", character_name=context.character_name
|
||||
)
|
||||
character.cover_image = asset.id
|
||||
|
||||
scene.emit_status()
|
||||
self.websocket_handler.request_scene_assets([asset.id])
|
||||
|
||||
self.websocket_handler.queue_put(
|
||||
{
|
||||
"type": "scene_asset_character_cover_image",
|
||||
"asset_id": asset.id,
|
||||
"asset": self.scene.assets.get_asset_bytes_as_base64(asset.id),
|
||||
"media_type": asset.media_type,
|
||||
"character": character.name,
|
||||
}
|
||||
)
|
||||
return
|
||||
@@ -1,42 +1,54 @@
|
||||
from __future__ import annotations
|
||||
import dataclasses
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import isodate
|
||||
import structlog
|
||||
|
||||
import talemate.emit.async_signals
|
||||
import talemate.util as util
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import DirectorMessage, TimePassageMessage
|
||||
from talemate.emit import emit
|
||||
from talemate.events import GameLoopEvent
|
||||
from talemate.instance import get_agent
|
||||
from talemate.prompts import Prompt
|
||||
from talemate.scene_message import (
|
||||
DirectorMessage,
|
||||
ReinforcementMessage,
|
||||
TimePassageMessage,
|
||||
)
|
||||
from talemate.world_state import InsertionMode
|
||||
|
||||
from .base import Agent, set_processing, AgentAction, AgentActionConfig, AgentEmission
|
||||
from .base import Agent, AgentAction, AgentActionConfig, AgentEmission, set_processing
|
||||
from .registry import register
|
||||
|
||||
import structlog
|
||||
import isodate
|
||||
import time
|
||||
|
||||
|
||||
log = structlog.get_logger("talemate.agents.world_state")
|
||||
|
||||
talemate.emit.async_signals.register("agent.world_state.time")
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class WorldStateAgentEmission(AgentEmission):
|
||||
"""
|
||||
Emission class for world state agent
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TimePassageEmission(WorldStateAgentEmission):
|
||||
"""
|
||||
Emission class for time passage
|
||||
"""
|
||||
|
||||
duration: str
|
||||
narrative: str
|
||||
|
||||
human_duration: str = None
|
||||
|
||||
|
||||
@register()
|
||||
class WorldStateAgent(Agent):
|
||||
@@ -51,21 +63,57 @@ class WorldStateAgent(Agent):
|
||||
self.client = client
|
||||
self.is_enabled = True
|
||||
self.actions = {
|
||||
"update_world_state": AgentAction(enabled=True, label="Update world state", description="Will attempt to update the world state based on the current scene. Runs automatically after AI dialogue (n turns).", config={
|
||||
"turns": AgentActionConfig(type="number", label="Turns", description="Number of turns to wait before updating the world state.", value=5, min=1, max=100, step=1)
|
||||
}),
|
||||
"update_world_state": AgentAction(
|
||||
enabled=True,
|
||||
label="Update world state",
|
||||
description="Will attempt to update the world state based on the current scene. Runs automatically every N turns.",
|
||||
config={
|
||||
"turns": AgentActionConfig(
|
||||
type="number",
|
||||
label="Turns",
|
||||
description="Number of turns to wait before updating the world state.",
|
||||
value=5,
|
||||
min=1,
|
||||
max=100,
|
||||
step=1,
|
||||
)
|
||||
},
|
||||
),
|
||||
"update_reinforcements": AgentAction(
|
||||
enabled=True,
|
||||
label="Update state reinforcements",
|
||||
description="Will attempt to update any due state reinforcements.",
|
||||
config={},
|
||||
),
|
||||
"check_pin_conditions": AgentAction(
|
||||
enabled=True,
|
||||
label="Update conditional context pins",
|
||||
description="Will evaluate context pins conditions and toggle those pins accordingly. Runs automatically every N turns.",
|
||||
config={
|
||||
"turns": AgentActionConfig(
|
||||
type="number",
|
||||
label="Turns",
|
||||
description="Number of turns to wait before checking conditions.",
|
||||
value=2,
|
||||
min=1,
|
||||
max=100,
|
||||
step=1,
|
||||
)
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
self.next_update = 0
|
||||
|
||||
self.next_pin_check = 0
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return self.is_enabled
|
||||
|
||||
|
||||
@property
|
||||
def has_toggle(self):
|
||||
return True
|
||||
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return True
|
||||
@@ -74,81 +122,121 @@ class WorldStateAgent(Agent):
|
||||
super().connect(scene)
|
||||
talemate.emit.async_signals.get("game_loop").connect(self.on_game_loop)
|
||||
|
||||
async def advance_time(self, duration:str, narrative:str=None):
|
||||
async def advance_time(self, duration: str, narrative: str = None):
|
||||
"""
|
||||
Emit a time passage message
|
||||
"""
|
||||
|
||||
|
||||
isodate.parse_duration(duration)
|
||||
msg_text = narrative or util.iso8601_duration_to_human(duration, suffix=" later")
|
||||
message = TimePassageMessage(ts=duration, message=msg_text)
|
||||
|
||||
human_duration = util.iso8601_duration_to_human(duration, suffix=" later")
|
||||
message = TimePassageMessage(ts=duration, message=human_duration)
|
||||
|
||||
log.debug("world_state.advance_time", message=message)
|
||||
self.scene.push_history(message)
|
||||
self.scene.emit_status()
|
||||
|
||||
emit("time", message)
|
||||
|
||||
await talemate.emit.async_signals.get("agent.world_state.time").send(
|
||||
TimePassageEmission(agent=self, duration=duration, narrative=msg_text)
|
||||
)
|
||||
|
||||
|
||||
async def on_game_loop(self, emission:GameLoopEvent):
|
||||
emit("time", message)
|
||||
|
||||
await talemate.emit.async_signals.get("agent.world_state.time").send(
|
||||
TimePassageEmission(
|
||||
agent=self,
|
||||
duration=duration,
|
||||
narrative=narrative,
|
||||
human_duration=human_duration,
|
||||
)
|
||||
)
|
||||
|
||||
async def on_game_loop(self, emission: GameLoopEvent):
|
||||
"""
|
||||
Called when a conversation is generated
|
||||
"""
|
||||
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
await self.update_world_state()
|
||||
|
||||
|
||||
async def update_world_state(self):
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
|
||||
await self.update_world_state()
|
||||
await self.auto_update_reinforcments()
|
||||
await self.auto_check_pin_conditions()
|
||||
|
||||
async def auto_update_reinforcments(self):
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
if not self.actions["update_reinforcements"].enabled:
|
||||
return
|
||||
|
||||
await self.update_reinforcements()
|
||||
|
||||
async def auto_check_pin_conditions(self):
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
if not self.actions["check_pin_conditions"].enabled:
|
||||
return
|
||||
|
||||
if (
|
||||
self.next_pin_check
|
||||
% self.actions["check_pin_conditions"].config["turns"].value
|
||||
!= 0
|
||||
or self.next_pin_check == 0
|
||||
):
|
||||
self.next_pin_check += 1
|
||||
return
|
||||
|
||||
self.next_pin_check = 0
|
||||
|
||||
await self.check_pin_conditions()
|
||||
|
||||
async def update_world_state(self, force: bool = False):
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
if not self.actions["update_world_state"].enabled:
|
||||
return
|
||||
|
||||
log.debug("update_world_state", next_update=self.next_update, turns=self.actions["update_world_state"].config["turns"].value)
|
||||
|
||||
|
||||
log.debug(
|
||||
"update_world_state",
|
||||
next_update=self.next_update,
|
||||
turns=self.actions["update_world_state"].config["turns"].value,
|
||||
)
|
||||
|
||||
scene = self.scene
|
||||
|
||||
if self.next_update % self.actions["update_world_state"].config["turns"].value != 0 or self.next_update == 0:
|
||||
|
||||
if (
|
||||
self.next_update % self.actions["update_world_state"].config["turns"].value
|
||||
!= 0
|
||||
or self.next_update == 0
|
||||
) and not force:
|
||||
self.next_update += 1
|
||||
return
|
||||
|
||||
|
||||
self.next_update = 0
|
||||
await scene.world_state.request_update()
|
||||
|
||||
|
||||
@set_processing
|
||||
async def request_world_state(self):
|
||||
|
||||
t1 = time.time()
|
||||
|
||||
_, world_state = await Prompt.request(
|
||||
"world_state.request-world-state-v2",
|
||||
self.client,
|
||||
"analyze_long",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"object_type": "character",
|
||||
"object_type_plural": "characters",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
self.scene.log.debug("request_world_state", response=world_state, time=time.time() - t1)
|
||||
|
||||
|
||||
self.scene.log.debug(
|
||||
"request_world_state", response=world_state, time=time.time() - t1
|
||||
)
|
||||
|
||||
return world_state
|
||||
|
||||
|
||||
@set_processing
|
||||
async def request_world_state_inline(self):
|
||||
|
||||
"""
|
||||
EXPERIMENTAL, Overall the one shot request seems about as coherent as the inline request, but the inline request is is about twice as slow and would need to run on every dialogue line.
|
||||
"""
|
||||
@@ -161,14 +249,18 @@ class WorldStateAgent(Agent):
|
||||
"world_state.request-world-state-inline-items",
|
||||
self.client,
|
||||
"analyze_long",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
self.scene.log.debug("request_world_state_inline", marked_items=marked_items_response, time=time.time() - t1)
|
||||
|
||||
|
||||
self.scene.log.debug(
|
||||
"request_world_state_inline",
|
||||
marked_items=marked_items_response,
|
||||
time=time.time() - t1,
|
||||
)
|
||||
|
||||
return marked_items_response
|
||||
|
||||
@set_processing
|
||||
@@ -176,70 +268,111 @@ class WorldStateAgent(Agent):
|
||||
self,
|
||||
text: str,
|
||||
):
|
||||
|
||||
response = await Prompt.request(
|
||||
"world_state.analyze-time-passage",
|
||||
self.client,
|
||||
"analyze_freeform_short",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
duration = response.split("\n")[0].split(" ")[0].strip()
|
||||
|
||||
|
||||
if not duration.startswith("P"):
|
||||
duration = "P"+duration
|
||||
|
||||
duration = "P" + duration
|
||||
|
||||
return duration
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def analyze_text_and_extract_context(
|
||||
self,
|
||||
text: str,
|
||||
goal: str,
|
||||
):
|
||||
|
||||
response = await Prompt.request(
|
||||
"world_state.analyze-text-and-extract-context",
|
||||
self.client,
|
||||
"analyze_freeform",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
"goal": goal,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
log.debug("analyze_text_and_extract_context", goal=goal, text=text, response=response)
|
||||
|
||||
|
||||
log.debug(
|
||||
"analyze_text_and_extract_context", goal=goal, text=text, response=response
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@set_processing
|
||||
async def analyze_text_and_extract_context_via_queries(
|
||||
self,
|
||||
text: str,
|
||||
goal: str,
|
||||
) -> list[str]:
|
||||
response = await Prompt.request(
|
||||
"world_state.analyze-text-and-generate-rag-queries",
|
||||
self.client,
|
||||
"analyze_freeform",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
"goal": goal,
|
||||
},
|
||||
)
|
||||
|
||||
queries = response.split("\n")
|
||||
|
||||
memory_agent = get_agent("memory")
|
||||
|
||||
context = await memory_agent.multi_query(queries, iterate=3)
|
||||
|
||||
log.debug(
|
||||
"analyze_text_and_extract_context_via_queries",
|
||||
goal=goal,
|
||||
text=text,
|
||||
queries=queries,
|
||||
context=context,
|
||||
)
|
||||
|
||||
return context
|
||||
|
||||
@set_processing
|
||||
async def analyze_and_follow_instruction(
|
||||
self,
|
||||
text: str,
|
||||
instruction: str,
|
||||
short: bool = False,
|
||||
):
|
||||
|
||||
|
||||
kind = "analyze_freeform_short" if short else "analyze_freeform"
|
||||
|
||||
response = await Prompt.request(
|
||||
"world_state.analyze-text-and-follow-instruction",
|
||||
self.client,
|
||||
"analyze_freeform",
|
||||
vars = {
|
||||
kind,
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
"instruction": instruction,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
log.debug("analyze_and_follow_instruction", instruction=instruction, text=text, response=response)
|
||||
|
||||
|
||||
log.debug(
|
||||
"analyze_and_follow_instruction",
|
||||
instruction=instruction,
|
||||
text=text,
|
||||
response=response,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
@@ -247,77 +380,55 @@ class WorldStateAgent(Agent):
|
||||
self,
|
||||
text: str,
|
||||
query: str,
|
||||
short: bool = False,
|
||||
):
|
||||
|
||||
kind = "analyze_freeform_short" if short else "analyze_freeform"
|
||||
response = await Prompt.request(
|
||||
"world_state.analyze-text-and-answer-question",
|
||||
self.client,
|
||||
"analyze_freeform",
|
||||
vars = {
|
||||
kind,
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
"query": query,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
log.debug("analyze_text_and_answer_question", query=query, text=text, response=response)
|
||||
|
||||
|
||||
log.debug(
|
||||
"analyze_text_and_answer_question",
|
||||
query=query,
|
||||
text=text,
|
||||
response=response,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@set_processing
|
||||
async def identify_characters(
|
||||
self,
|
||||
text: str = None,
|
||||
):
|
||||
|
||||
"""
|
||||
Attempts to identify characters in the given text.
|
||||
"""
|
||||
|
||||
|
||||
_, data = await Prompt.request(
|
||||
"world_state.identify-characters",
|
||||
self.client,
|
||||
"analyze",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
log.debug("identify_characters", text=text, data=data)
|
||||
|
||||
|
||||
return data
|
||||
|
||||
@set_processing
|
||||
async def extract_character_sheet(
|
||||
self,
|
||||
name:str,
|
||||
text:str = None,
|
||||
):
|
||||
|
||||
"""
|
||||
Attempts to extract a character sheet from the given text.
|
||||
"""
|
||||
|
||||
response = await Prompt.request(
|
||||
"world_state.extract-character-sheet",
|
||||
self.client,
|
||||
"analyze_creative",
|
||||
vars = {
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
|
||||
# loop through each line in response and if it contains a : then extract
|
||||
# the left side as an attribute name and the right side as the value
|
||||
#
|
||||
# break as soon as a non-empty line is found that doesn't contain a :
|
||||
|
||||
|
||||
def _parse_character_sheet(self, response):
|
||||
data = {}
|
||||
for line in response.split("\n"):
|
||||
if not line.strip():
|
||||
@@ -326,28 +437,349 @@ class WorldStateAgent(Agent):
|
||||
break
|
||||
name, value = line.split(":", 1)
|
||||
data[name.strip()] = value.strip()
|
||||
|
||||
|
||||
return data
|
||||
|
||||
|
||||
|
||||
@set_processing
|
||||
async def match_character_names(self, names:list[str]):
|
||||
|
||||
async def extract_character_sheet(
|
||||
self,
|
||||
name: str,
|
||||
text: str = None,
|
||||
alteration_instructions: str = None,
|
||||
):
|
||||
"""
|
||||
Attempts to extract a character sheet from the given text.
|
||||
"""
|
||||
|
||||
response = await Prompt.request(
|
||||
"world_state.extract-character-sheet",
|
||||
self.client,
|
||||
"create",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"text": text,
|
||||
"name": name,
|
||||
"character": self.scene.get_character(name),
|
||||
"alteration_instructions": alteration_instructions or "",
|
||||
},
|
||||
)
|
||||
|
||||
# loop through each line in response and if it contains a : then extract
|
||||
# the left side as an attribute name and the right side as the value
|
||||
#
|
||||
# break as soon as a non-empty line is found that doesn't contain a :
|
||||
|
||||
return self._parse_character_sheet(response)
|
||||
|
||||
@set_processing
|
||||
async def match_character_names(self, names: list[str]):
|
||||
"""
|
||||
Attempts to match character names.
|
||||
"""
|
||||
|
||||
|
||||
_, response = await Prompt.request(
|
||||
"world_state.match-character-names",
|
||||
self.client,
|
||||
"analyze_long",
|
||||
vars = {
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"names": names,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
log.debug("match_character_names", names=names, response=response)
|
||||
|
||||
return response
|
||||
|
||||
return response
|
||||
|
||||
@set_processing
|
||||
async def update_reinforcements(self, force: bool = False):
|
||||
"""
|
||||
Queries due worldstate re-inforcements
|
||||
"""
|
||||
|
||||
for reinforcement in self.scene.world_state.reinforce:
|
||||
if reinforcement.due <= 0 or force:
|
||||
await self.update_reinforcement(
|
||||
reinforcement.question, reinforcement.character
|
||||
)
|
||||
else:
|
||||
reinforcement.due -= 1
|
||||
|
||||
@set_processing
|
||||
async def update_reinforcement(
|
||||
self, question: str, character: str = None, reset: bool = False
|
||||
):
|
||||
"""
|
||||
Queries a single re-inforcement
|
||||
"""
|
||||
message = None
|
||||
idx, reinforcement = await self.scene.world_state.find_reinforcement(
|
||||
question, character
|
||||
)
|
||||
|
||||
if not reinforcement:
|
||||
return
|
||||
|
||||
source = f"{reinforcement.question}:{reinforcement.character if reinforcement.character else ''}"
|
||||
|
||||
if reset and reinforcement.insert == "sequential":
|
||||
self.scene.pop_history(typ="reinforcement", source=source, all=True)
|
||||
|
||||
if reinforcement.insert == "sequential":
|
||||
kind = "analyze_freeform_medium_short"
|
||||
else:
|
||||
kind = "analyze_freeform"
|
||||
|
||||
answer = await Prompt.request(
|
||||
"world_state.update-reinforcements",
|
||||
self.client,
|
||||
kind,
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"question": reinforcement.question,
|
||||
"instructions": reinforcement.instructions or "",
|
||||
"character": (
|
||||
self.scene.get_character(reinforcement.character)
|
||||
if reinforcement.character
|
||||
else None
|
||||
),
|
||||
"answer": (reinforcement.answer if not reset else None) or "",
|
||||
"reinforcement": reinforcement,
|
||||
},
|
||||
)
|
||||
|
||||
# sequential reinforcment should be single sentence so we
|
||||
# split on line breaks and take the first line in case the
|
||||
# LLM did not understand the request and returned a longer response
|
||||
|
||||
if reinforcement.insert == "sequential":
|
||||
answer = answer.split("\n")[0]
|
||||
|
||||
reinforcement.answer = answer
|
||||
reinforcement.due = reinforcement.interval
|
||||
|
||||
# remove any recent previous reinforcement message with same question
|
||||
# to avoid overloading the near history with reinforcement messages
|
||||
if not reset:
|
||||
self.scene.pop_history(
|
||||
typ="reinforcement", source=source, max_iterations=10
|
||||
)
|
||||
|
||||
if reinforcement.insert == "sequential":
|
||||
# insert the reinforcement message at the current position
|
||||
message = ReinforcementMessage(message=answer, source=source)
|
||||
log.debug("update_reinforcement", message=message, reset=reset)
|
||||
self.scene.push_history(message)
|
||||
|
||||
# if reinforcement has a character name set, update the character detail
|
||||
if reinforcement.character:
|
||||
character = self.scene.get_character(reinforcement.character)
|
||||
await character.set_detail(reinforcement.question, answer)
|
||||
|
||||
else:
|
||||
# set world entry
|
||||
await self.scene.world_state_manager.save_world_entry(
|
||||
reinforcement.question,
|
||||
reinforcement.as_context_line,
|
||||
{},
|
||||
)
|
||||
|
||||
self.scene.world_state.emit()
|
||||
|
||||
return message
|
||||
|
||||
@set_processing
|
||||
async def check_pin_conditions(
|
||||
self,
|
||||
):
|
||||
"""
|
||||
Checks if any context pin conditions
|
||||
"""
|
||||
|
||||
pins_with_condition = {
|
||||
entry_id: {
|
||||
"condition": pin.condition,
|
||||
"state": pin.condition_state,
|
||||
}
|
||||
for entry_id, pin in self.scene.world_state.pins.items()
|
||||
if pin.condition
|
||||
}
|
||||
|
||||
if not pins_with_condition:
|
||||
return
|
||||
|
||||
first_entry_id = list(pins_with_condition.keys())[0]
|
||||
|
||||
_, answers = await Prompt.request(
|
||||
"world_state.check-pin-conditions",
|
||||
self.client,
|
||||
"analyze",
|
||||
vars={
|
||||
"scene": self.scene,
|
||||
"max_tokens": self.client.max_token_length,
|
||||
"previous_states": json.dumps(pins_with_condition, indent=2),
|
||||
"coercion": {first_entry_id: {"condition": ""}},
|
||||
},
|
||||
)
|
||||
|
||||
world_state = self.scene.world_state
|
||||
state_change = False
|
||||
|
||||
for entry_id, answer in answers.items():
|
||||
if entry_id not in world_state.pins:
|
||||
log.warning(
|
||||
"check_pin_conditions",
|
||||
entry_id=entry_id,
|
||||
answer=answer,
|
||||
msg="entry_id not found in world_state.pins (LLM failed to produce a clean response)",
|
||||
)
|
||||
continue
|
||||
|
||||
log.info("check_pin_conditions", entry_id=entry_id, answer=answer)
|
||||
state = answer.get("state")
|
||||
if state is True or (
|
||||
isinstance(state, str) and state.lower() in ["true", "yes", "y"]
|
||||
):
|
||||
prev_state = world_state.pins[entry_id].condition_state
|
||||
|
||||
world_state.pins[entry_id].condition_state = True
|
||||
world_state.pins[entry_id].active = True
|
||||
|
||||
if prev_state != world_state.pins[entry_id].condition_state:
|
||||
state_change = True
|
||||
else:
|
||||
if world_state.pins[entry_id].condition_state is not False:
|
||||
world_state.pins[entry_id].condition_state = False
|
||||
world_state.pins[entry_id].active = False
|
||||
state_change = True
|
||||
|
||||
if state_change:
|
||||
await self.scene.load_active_pins()
|
||||
self.scene.emit_status()
|
||||
|
||||
@set_processing
|
||||
async def summarize_and_pin(self, message_id: int, num_messages: int = 3) -> str:
|
||||
"""
|
||||
Will take a message index and then walk back N messages
|
||||
summarizing the scene and pinning it to the context.
|
||||
"""
|
||||
|
||||
creator = get_agent("creator")
|
||||
summarizer = get_agent("summarizer")
|
||||
|
||||
message_index = self.scene.message_index(message_id)
|
||||
|
||||
text = self.scene.snapshot(lines=num_messages, start=message_index)
|
||||
|
||||
extra_context = self.scene.snapshot(
|
||||
lines=50, start=message_index - num_messages
|
||||
)
|
||||
|
||||
summary = await summarizer.summarize(
|
||||
text,
|
||||
extra_context=extra_context,
|
||||
method="short",
|
||||
extra_instructions="Pay particularly close attention to decisions, agreements or promises made.",
|
||||
)
|
||||
|
||||
entry_id = util.clean_id(await creator.generate_title(summary))
|
||||
|
||||
ts = self.scene.ts
|
||||
|
||||
log.debug(
|
||||
"summarize_and_pin",
|
||||
message_id=message_id,
|
||||
message_index=message_index,
|
||||
num_messages=num_messages,
|
||||
summary=summary,
|
||||
entry_id=entry_id,
|
||||
ts=ts,
|
||||
)
|
||||
|
||||
await self.scene.world_state_manager.save_world_entry(
|
||||
entry_id,
|
||||
summary,
|
||||
{
|
||||
"ts": ts,
|
||||
},
|
||||
)
|
||||
|
||||
await self.scene.world_state_manager.set_pin(
|
||||
entry_id,
|
||||
active=True,
|
||||
)
|
||||
|
||||
await self.scene.load_active_pins()
|
||||
self.scene.emit_status()
|
||||
|
||||
@set_processing
|
||||
async def is_character_present(self, character: str) -> bool:
|
||||
"""
|
||||
Check if a character is present in the scene
|
||||
|
||||
Arguments:
|
||||
|
||||
- `character`: The character to check.
|
||||
"""
|
||||
|
||||
if len(self.scene.history) < 10:
|
||||
text = self.scene.intro + "\n\n" + self.scene.snapshot(lines=50)
|
||||
else:
|
||||
text = self.scene.snapshot(lines=50)
|
||||
|
||||
is_present = await self.analyze_text_and_answer_question(
|
||||
text=text,
|
||||
query=f"Is {character} present AND active in the current scene? Answert with 'yes' or 'no'.",
|
||||
)
|
||||
|
||||
return is_present.lower().startswith("y")
|
||||
|
||||
@set_processing
|
||||
async def is_character_leaving(self, character: str) -> bool:
|
||||
"""
|
||||
Check if a character is leaving the scene
|
||||
|
||||
Arguments:
|
||||
|
||||
- `character`: The character to check.
|
||||
"""
|
||||
|
||||
if len(self.scene.history) < 10:
|
||||
text = self.scene.intro + "\n\n" + self.scene.snapshot(lines=50)
|
||||
else:
|
||||
text = self.scene.snapshot(lines=50)
|
||||
|
||||
is_leaving = await self.analyze_text_and_answer_question(
|
||||
text=text,
|
||||
query=f"Is {character} leaving the current scene? Answert with 'yes' or 'no'.",
|
||||
)
|
||||
|
||||
return is_leaving.lower().startswith("y")
|
||||
|
||||
@set_processing
|
||||
async def manager(self, action_name: str, *args, **kwargs):
|
||||
"""
|
||||
Executes a world state manager action through self.scene.world_state_manager
|
||||
"""
|
||||
|
||||
manager = self.scene.world_state_manager
|
||||
|
||||
try:
|
||||
fn = getattr(manager, action_name, None)
|
||||
|
||||
if not fn:
|
||||
raise ValueError(f"Unknown action: {action_name}")
|
||||
|
||||
return await fn(*args, **kwargs)
|
||||
except Exception as e:
|
||||
log.error(
|
||||
"worldstate.manager",
|
||||
action_name=action_name,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
error=e,
|
||||
)
|
||||
raise
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
from __future__ import annotations
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import dataclasses
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate import Scene
|
||||
|
||||
|
||||
import structlog
|
||||
|
||||
__all__ = ["AutomatedAction", "register", "initialize_for_scene"]
|
||||
@@ -13,50 +14,64 @@ log = structlog.get_logger("talemate.automated_action")
|
||||
|
||||
AUTOMATED_ACTIONS = {}
|
||||
|
||||
def initialize_for_scene(scene:Scene):
|
||||
|
||||
|
||||
def initialize_for_scene(scene: Scene):
|
||||
for uid, config in AUTOMATED_ACTIONS.items():
|
||||
scene.automated_actions[uid] = config.cls(
|
||||
scene,
|
||||
uid=uid,
|
||||
frequency=config.frequency,
|
||||
call_initially=config.call_initially,
|
||||
enabled=config.enabled
|
||||
enabled=config.enabled,
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AutomatedActionConfig:
|
||||
uid:str
|
||||
cls:AutomatedAction
|
||||
frequency:int=5
|
||||
call_initially:bool=False
|
||||
enabled:bool=True
|
||||
|
||||
uid: str
|
||||
cls: AutomatedAction
|
||||
frequency: int = 5
|
||||
call_initially: bool = False
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
class register:
|
||||
|
||||
def __init__(self, uid:str, frequency:int=5, call_initially:bool=False, enabled:bool=True):
|
||||
def __init__(
|
||||
self,
|
||||
uid: str,
|
||||
frequency: int = 5,
|
||||
call_initially: bool = False,
|
||||
enabled: bool = True,
|
||||
):
|
||||
self.uid = uid
|
||||
self.frequency = frequency
|
||||
self.call_initially = call_initially
|
||||
self.enabled = enabled
|
||||
|
||||
def __call__(self, action:AutomatedAction):
|
||||
|
||||
def __call__(self, action: AutomatedAction):
|
||||
AUTOMATED_ACTIONS[self.uid] = AutomatedActionConfig(
|
||||
self.uid,
|
||||
action,
|
||||
frequency=self.frequency,
|
||||
call_initially=self.call_initially,
|
||||
enabled=self.enabled
|
||||
self.uid,
|
||||
action,
|
||||
frequency=self.frequency,
|
||||
call_initially=self.call_initially,
|
||||
enabled=self.enabled,
|
||||
)
|
||||
return action
|
||||
|
||||
|
||||
|
||||
class AutomatedAction:
|
||||
"""
|
||||
An action that will be executed every n turns
|
||||
"""
|
||||
|
||||
def __init__(self, scene:Scene, frequency:int=5, call_initially:bool=False, uid:str=None, enabled:bool=True):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scene: Scene,
|
||||
frequency: int = 5,
|
||||
call_initially: bool = False,
|
||||
uid: str = None,
|
||||
enabled: bool = True,
|
||||
):
|
||||
self.scene = scene
|
||||
self.enabled = enabled
|
||||
self.frequency = frequency
|
||||
@@ -64,14 +79,19 @@ class AutomatedAction:
|
||||
self.uid = uid
|
||||
if call_initially:
|
||||
self.turns = frequency
|
||||
|
||||
|
||||
async def __call__(self):
|
||||
|
||||
log.debug("automated_action", uid=self.uid, enabled=self.enabled, frequency=self.frequency, turns=self.turns)
|
||||
|
||||
log.debug(
|
||||
"automated_action",
|
||||
uid=self.uid,
|
||||
enabled=self.enabled,
|
||||
frequency=self.frequency,
|
||||
turns=self.turns,
|
||||
)
|
||||
|
||||
if not self.enabled:
|
||||
return False
|
||||
|
||||
|
||||
if self.turns % self.frequency == 0:
|
||||
result = await self.action()
|
||||
log.debug("automated_action", result=result)
|
||||
@@ -79,10 +99,9 @@ class AutomatedAction:
|
||||
# action could not be performed at this turn, we will try again next turn
|
||||
return False
|
||||
self.turns += 1
|
||||
|
||||
|
||||
|
||||
async def action(self) -> Any:
|
||||
"""
|
||||
Override this method to implement your action.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
raise NotImplementedError()
|
||||
|
||||
59
src/talemate/character.py
Normal file
@@ -0,0 +1,59 @@
|
||||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
from talemate.instance import get_agent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from talemate.tale_mate import Actor, Character, Scene
|
||||
|
||||
|
||||
__all__ = [
|
||||
"deactivate_character",
|
||||
"activate_character",
|
||||
]
|
||||
|
||||
|
||||
async def deactivate_character(scene: "Scene", character: Union[str, "Character"]):
|
||||
"""
|
||||
Deactivates a character
|
||||
|
||||
Arguments:
|
||||
|
||||
- `scene`: The scene to deactivate the character from
|
||||
- `character`: The character to deactivate. Can be a string (the character's name) or a Character object
|
||||
"""
|
||||
|
||||
if isinstance(character, str):
|
||||
character = scene.get_character(character)
|
||||
|
||||
if character.is_player:
|
||||
# can't deactivate the player
|
||||
return False
|
||||
|
||||
if character.name in scene.inactive_characters:
|
||||
# already deactivated
|
||||
return False
|
||||
|
||||
await scene.remove_actor(character.actor)
|
||||
scene.inactive_characters[character.name] = character
|
||||
|
||||
|
||||
async def activate_character(scene: "Scene", character: Union[str, "Character"]):
|
||||
"""
|
||||
Activates a character
|
||||
|
||||
Arguments:
|
||||
|
||||
- `scene`: The scene to activate the character in
|
||||
- `character`: The character to activate. Can be a string (the character's name) or a Character object
|
||||
"""
|
||||
|
||||
if isinstance(character, str):
|
||||
character = scene.get_character(character)
|
||||
|
||||
if character.name not in scene.inactive_characters:
|
||||
# already activated
|
||||
return False
|
||||
|
||||
actor = scene.Actor(character, get_agent("conversation"))
|
||||
await scene.add_actor(actor)
|
||||
del scene.inactive_characters[character.name]
|
||||
@@ -2,15 +2,13 @@ import argparse
|
||||
import asyncio
|
||||
import glob
|
||||
import os
|
||||
import structlog
|
||||
|
||||
import structlog
|
||||
from dotenv import load_dotenv
|
||||
|
||||
import talemate.instance as instance
|
||||
from talemate import Actor, Character, Helper, Player, Scene
|
||||
from talemate.agents import (
|
||||
ConversationAgent,
|
||||
)
|
||||
from talemate.agents import ConversationAgent
|
||||
from talemate.client import OpenAIClient, TextGeneratorWebuiClient
|
||||
from talemate.emit.console import Console
|
||||
from talemate.load import (
|
||||
@@ -129,7 +127,6 @@ async def run_console_session(parser, args):
|
||||
default_client = None
|
||||
|
||||
if "textgenwebui" in clients.values() or args.client == "textgenwebui":
|
||||
|
||||
# Init the TextGeneratorWebuiClient with ConversationAgent and create an actor
|
||||
textgenwebui_api_url = args.textgenwebui_url
|
||||
|
||||
@@ -145,7 +142,6 @@ async def run_console_session(parser, args):
|
||||
clients[client_name] = text_generator_webui_client
|
||||
|
||||
if "openai" in clients.values() or args.client == "openai":
|
||||
|
||||
openai_client = OpenAIClient()
|
||||
|
||||
for client_name, client_typ in clients.items():
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
import os
|
||||
|
||||
import talemate.client.runpod
|
||||
from talemate.client.lmstudio import LMStudioClient
|
||||
from talemate.client.openai import OpenAIClient
|
||||
from talemate.client.mistral import MistralAIClient
|
||||
from talemate.client.anthropic import AnthropicClient
|
||||
from talemate.client.openai_compat import OpenAICompatibleClient
|
||||
from talemate.client.registry import CLIENT_CLASSES, get_client_class, register
|
||||
from talemate.client.textgenwebui import TextGeneratorWebuiClient
|
||||
from talemate.client.lmstudio import LMStudioClient
|
||||
import talemate.client.runpod
|
||||
|
||||
224
src/talemate/client/anthropic.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import pydantic
|
||||
import structlog
|
||||
from anthropic import AsyncAnthropic, PermissionDeniedError
|
||||
|
||||
from talemate.client.base import ClientBase, ErrorAction
|
||||
from talemate.client.registry import register
|
||||
from talemate.config import load_config
|
||||
from talemate.emit import emit
|
||||
from talemate.emit.signals import handlers
|
||||
|
||||
__all__ = [
|
||||
"AnthropicClient",
|
||||
]
|
||||
log = structlog.get_logger("talemate")
|
||||
|
||||
# Edit this to add new models / remove old models
|
||||
SUPPORTED_MODELS = [
|
||||
"claude-3-sonnet-20240229",
|
||||
"claude-3-opus-20240229",
|
||||
]
|
||||
|
||||
|
||||
class Defaults(pydantic.BaseModel):
|
||||
max_token_length: int = 16384
|
||||
model: str = "claude-3-sonnet-20240229"
|
||||
|
||||
|
||||
@register()
|
||||
class AnthropicClient(ClientBase):
|
||||
"""
|
||||
Anthropic client for generating text.
|
||||
"""
|
||||
|
||||
client_type = "anthropic"
|
||||
conversation_retries = 0
|
||||
auto_break_repetition_enabled = False
|
||||
# TODO: make this configurable?
|
||||
decensor_enabled = False
|
||||
|
||||
class Meta(ClientBase.Meta):
|
||||
name_prefix: str = "Anthropic"
|
||||
title: str = "Anthropic"
|
||||
manual_model: bool = True
|
||||
manual_model_choices: list[str] = SUPPORTED_MODELS
|
||||
requires_prompt_template: bool = False
|
||||
defaults: Defaults = Defaults()
|
||||
|
||||
def __init__(self, model="claude-3-sonnet-20240229", **kwargs):
|
||||
self.model_name = model
|
||||
self.api_key_status = None
|
||||
self.config = load_config()
|
||||
super().__init__(**kwargs)
|
||||
|
||||
handlers["config_saved"].connect(self.on_config_saved)
|
||||
|
||||
@property
|
||||
def anthropic_api_key(self):
|
||||
return self.config.get("anthropic", {}).get("api_key")
|
||||
|
||||
def emit_status(self, processing: bool = None):
|
||||
error_action = None
|
||||
if processing is not None:
|
||||
self.processing = processing
|
||||
|
||||
if self.anthropic_api_key:
|
||||
status = "busy" if self.processing else "idle"
|
||||
model_name = self.model_name
|
||||
else:
|
||||
status = "error"
|
||||
model_name = "No API key set"
|
||||
error_action = ErrorAction(
|
||||
title="Set API Key",
|
||||
action_name="openAppConfig",
|
||||
icon="mdi-key-variant",
|
||||
arguments=[
|
||||
"application",
|
||||
"anthropic_api",
|
||||
],
|
||||
)
|
||||
|
||||
if not self.model_name:
|
||||
status = "error"
|
||||
model_name = "No model loaded"
|
||||
|
||||
self.current_status = status
|
||||
|
||||
emit(
|
||||
"client_status",
|
||||
message=self.client_type,
|
||||
id=self.name,
|
||||
details=model_name,
|
||||
status=status,
|
||||
data={
|
||||
"error_action": error_action.model_dump() if error_action else None,
|
||||
"meta": self.Meta().model_dump(),
|
||||
},
|
||||
)
|
||||
|
||||
def set_client(self, max_token_length: int = None):
|
||||
if not self.anthropic_api_key:
|
||||
self.client = AsyncAnthropic(api_key="sk-1111")
|
||||
log.error("No anthropic API key set")
|
||||
if self.api_key_status:
|
||||
self.api_key_status = False
|
||||
emit("request_client_status")
|
||||
emit("request_agent_status")
|
||||
return
|
||||
|
||||
if not self.model_name:
|
||||
self.model_name = "claude-3-opus-20240229"
|
||||
|
||||
if max_token_length and not isinstance(max_token_length, int):
|
||||
max_token_length = int(max_token_length)
|
||||
|
||||
model = self.model_name
|
||||
|
||||
self.client = AsyncAnthropic(api_key=self.anthropic_api_key)
|
||||
self.max_token_length = max_token_length or 16384
|
||||
|
||||
if not self.api_key_status:
|
||||
if self.api_key_status is False:
|
||||
emit("request_client_status")
|
||||
emit("request_agent_status")
|
||||
self.api_key_status = True
|
||||
|
||||
log.info(
|
||||
"anthropic set client",
|
||||
max_token_length=self.max_token_length,
|
||||
provided_max_token_length=max_token_length,
|
||||
model=model,
|
||||
)
|
||||
|
||||
def reconfigure(self, **kwargs):
|
||||
if kwargs.get("model"):
|
||||
self.model_name = kwargs["model"]
|
||||
self.set_client(kwargs.get("max_token_length"))
|
||||
|
||||
def on_config_saved(self, event):
|
||||
config = event.data
|
||||
self.config = config
|
||||
self.set_client(max_token_length=self.max_token_length)
|
||||
|
||||
def response_tokens(self, response: str):
|
||||
return response.usage.output_tokens
|
||||
|
||||
def prompt_tokens(self, response: str):
|
||||
return response.usage.input_tokens
|
||||
|
||||
async def status(self):
|
||||
self.emit_status()
|
||||
|
||||
def prompt_template(self, system_message: str, prompt: str):
|
||||
if "<|BOT|>" in prompt:
|
||||
_, right = prompt.split("<|BOT|>", 1)
|
||||
if right:
|
||||
prompt = prompt.replace("<|BOT|>", "\nStart your response with: ")
|
||||
else:
|
||||
prompt = prompt.replace("<|BOT|>", "")
|
||||
|
||||
return prompt
|
||||
|
||||
def tune_prompt_parameters(self, parameters: dict, kind: str):
|
||||
super().tune_prompt_parameters(parameters, kind)
|
||||
keys = list(parameters.keys())
|
||||
valid_keys = ["temperature", "top_p", "max_tokens"]
|
||||
for key in keys:
|
||||
if key not in valid_keys:
|
||||
del parameters[key]
|
||||
|
||||
async def generate(self, prompt: str, parameters: dict, kind: str):
|
||||
"""
|
||||
Generates text from the given prompt and parameters.
|
||||
"""
|
||||
|
||||
if not self.anthropic_api_key:
|
||||
raise Exception("No anthropic API key set")
|
||||
|
||||
right = None
|
||||
expected_response = None
|
||||
try:
|
||||
_, right = prompt.split("\nStart your response with: ")
|
||||
expected_response = right.strip()
|
||||
except (IndexError, ValueError):
|
||||
pass
|
||||
|
||||
human_message = {"role": "user", "content": prompt.strip()}
|
||||
system_message = self.get_system_message(kind)
|
||||
|
||||
self.log.debug(
|
||||
"generate",
|
||||
prompt=prompt[:128] + " ...",
|
||||
parameters=parameters,
|
||||
system_message=system_message,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self.client.messages.create(
|
||||
model=self.model_name,
|
||||
system=system_message,
|
||||
messages=[human_message],
|
||||
**parameters,
|
||||
)
|
||||
|
||||
self._returned_prompt_tokens = self.prompt_tokens(response)
|
||||
self._returned_response_tokens = self.response_tokens(response)
|
||||
|
||||
log.debug("generated response", response=response.content)
|
||||
|
||||
response = response.content[0].text
|
||||
|
||||
if expected_response and expected_response.startswith("{"):
|
||||
if response.startswith("```json") and response.endswith("```"):
|
||||
response = response[7:-3].strip()
|
||||
|
||||
if right and response.startswith(right):
|
||||
response = response[len(right) :].strip()
|
||||
|
||||
return response
|
||||
except PermissionDeniedError as e:
|
||||
self.log.error("generate error", e=e)
|
||||
emit("status", message="anthropic API: Permission Denied", status="error")
|
||||
return ""
|
||||
except Exception as e:
|
||||
raise
|
||||
@@ -1,26 +1,29 @@
|
||||
"""
|
||||
A unified client base, based on the openai API
|
||||
"""
|
||||
import copy
|
||||
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from typing import Callable
|
||||
from typing import Callable, Union
|
||||
|
||||
import pydantic
|
||||
import structlog
|
||||
import logging
|
||||
from openai import AsyncOpenAI
|
||||
from openai import AsyncOpenAI, PermissionDeniedError
|
||||
|
||||
from talemate.emit import emit
|
||||
import talemate.instance as instance
|
||||
import talemate.client.presets as presets
|
||||
import talemate.client.system_prompts as system_prompts
|
||||
import talemate.instance as instance
|
||||
import talemate.util as util
|
||||
from talemate.agents.context import active_agent
|
||||
from talemate.client.context import client_context_attribute
|
||||
from talemate.client.model_prompts import model_prompt
|
||||
|
||||
from talemate.emit import emit
|
||||
|
||||
# Set up logging level for httpx to WARNING to suppress debug logs.
|
||||
logging.getLogger('httpx').setLevel(logging.WARNING)
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
log = structlog.get_logger("client.base")
|
||||
|
||||
REMOTE_SERVICES = [
|
||||
# TODO: runpod.py should add this to the list
|
||||
@@ -29,133 +32,216 @@ REMOTE_SERVICES = [
|
||||
|
||||
STOPPING_STRINGS = ["<|im_end|>", "</s>"]
|
||||
|
||||
|
||||
class PromptData(pydantic.BaseModel):
|
||||
kind: str
|
||||
prompt: str
|
||||
response: str
|
||||
prompt_tokens: int
|
||||
response_tokens: int
|
||||
client_name: str
|
||||
client_type: str
|
||||
time: Union[float, int]
|
||||
agent_stack: list[str] = pydantic.Field(default_factory=list)
|
||||
generation_parameters: dict = pydantic.Field(default_factory=dict)
|
||||
|
||||
|
||||
class ErrorAction(pydantic.BaseModel):
|
||||
title: str
|
||||
action_name: str
|
||||
icon: str = "mdi-error"
|
||||
arguments: list = []
|
||||
|
||||
|
||||
class Defaults(pydantic.BaseModel):
|
||||
api_url: str = "http://localhost:5000"
|
||||
max_token_length: int = 4096
|
||||
|
||||
|
||||
class ExtraField(pydantic.BaseModel):
|
||||
name: str
|
||||
type: str
|
||||
label: str
|
||||
required: bool
|
||||
description: str
|
||||
|
||||
|
||||
class ClientBase:
|
||||
|
||||
api_url: str
|
||||
model_name: str
|
||||
name:str = None
|
||||
api_key: str = None
|
||||
name: str = None
|
||||
enabled: bool = True
|
||||
current_status: str = None
|
||||
max_token_length: int = 4096
|
||||
randomizable_inference_parameters: list[str] = ["temperature"]
|
||||
processing: bool = False
|
||||
connected: bool = False
|
||||
conversation_retries: int = 5
|
||||
|
||||
conversation_retries: int = 2
|
||||
auto_break_repetition_enabled: bool = True
|
||||
decensor_enabled: bool = True
|
||||
client_type = "base"
|
||||
|
||||
|
||||
class Meta(pydantic.BaseModel):
|
||||
experimental: Union[None, str] = None
|
||||
defaults: Defaults = Defaults()
|
||||
title: str = "Client"
|
||||
name_prefix: str = "Client"
|
||||
enable_api_auth: bool = False
|
||||
requires_prompt_template: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_url: str = None,
|
||||
name = None,
|
||||
name=None,
|
||||
**kwargs,
|
||||
):
|
||||
self.api_url = api_url
|
||||
self.name = name or self.client_type
|
||||
self.log = structlog.get_logger(f"client.{self.client_type}")
|
||||
self.set_client()
|
||||
|
||||
if "max_token_length" in kwargs:
|
||||
self.max_token_length = (
|
||||
int(kwargs["max_token_length"]) if kwargs["max_token_length"] else 4096
|
||||
)
|
||||
self.set_client(max_token_length=self.max_token_length)
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.client_type}Client[{self.api_url}][{self.model_name or ''}]"
|
||||
|
||||
def set_client(self):
|
||||
|
||||
@property
|
||||
def experimental(self):
|
||||
return False
|
||||
|
||||
def set_client(self, **kwargs):
|
||||
self.client = AsyncOpenAI(base_url=self.api_url, api_key="sk-1111")
|
||||
|
||||
def prompt_template(self, sys_msg, prompt):
|
||||
|
||||
"""
|
||||
Applies the appropriate prompt template for the model.
|
||||
"""
|
||||
|
||||
|
||||
if not self.model_name:
|
||||
self.log.warning("prompt template not applied", reason="no model loaded")
|
||||
return f"{sys_msg}\n{prompt}"
|
||||
|
||||
return model_prompt(self.model_name, sys_msg, prompt)
|
||||
|
||||
|
||||
return model_prompt(self.model_name, sys_msg, prompt)[0]
|
||||
|
||||
def prompt_template_example(self):
|
||||
if not getattr(self, "model_name", None):
|
||||
return None, None
|
||||
return model_prompt(self.model_name, "sysmsg", "prompt<|BOT|>{LLM coercion}")
|
||||
|
||||
def reconfigure(self, **kwargs):
|
||||
|
||||
"""
|
||||
Reconfigures the client.
|
||||
|
||||
|
||||
Keyword Arguments:
|
||||
|
||||
|
||||
- api_url: the API URL to use
|
||||
- max_token_length: the max token length to use
|
||||
- enabled: whether the client is enabled
|
||||
"""
|
||||
|
||||
|
||||
if "api_url" in kwargs:
|
||||
self.api_url = kwargs["api_url"]
|
||||
|
||||
if "max_token_length" in kwargs:
|
||||
self.max_token_length = kwargs["max_token_length"]
|
||||
|
||||
if kwargs.get("max_token_length"):
|
||||
self.max_token_length = int(kwargs["max_token_length"])
|
||||
|
||||
if "enabled" in kwargs:
|
||||
self.enabled = bool(kwargs["enabled"])
|
||||
|
||||
|
||||
def toggle_disabled_if_remote(self):
|
||||
|
||||
"""
|
||||
If the client is targeting a remote recognized service, this
|
||||
will disable the client.
|
||||
"""
|
||||
|
||||
|
||||
for service in REMOTE_SERVICES:
|
||||
if service in self.api_url:
|
||||
if self.enabled:
|
||||
self.log.warn("remote service unreachable, disabling client", client=self.name)
|
||||
self.log.warn(
|
||||
"remote service unreachable, disabling client", client=self.name
|
||||
)
|
||||
self.enabled = False
|
||||
|
||||
|
||||
return True
|
||||
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def get_system_message(self, kind: str) -> str:
|
||||
|
||||
"""
|
||||
Returns the appropriate system message for the given kind of generation
|
||||
|
||||
|
||||
Arguments:
|
||||
|
||||
|
||||
- kind: the kind of generation
|
||||
"""
|
||||
|
||||
# TODO: make extensible
|
||||
|
||||
if "narrate" in kind:
|
||||
return system_prompts.NARRATOR
|
||||
if "story" in kind:
|
||||
return system_prompts.NARRATOR
|
||||
if "director" in kind:
|
||||
return system_prompts.DIRECTOR
|
||||
if "create" in kind:
|
||||
return system_prompts.CREATOR
|
||||
if "roleplay" in kind:
|
||||
return system_prompts.ROLEPLAY
|
||||
if "conversation" in kind:
|
||||
return system_prompts.ROLEPLAY
|
||||
if "editor" in kind:
|
||||
return system_prompts.EDITOR
|
||||
if "world_state" in kind:
|
||||
return system_prompts.WORLD_STATE
|
||||
if "analyst" in kind:
|
||||
return system_prompts.ANALYST
|
||||
if "analyze" in kind:
|
||||
return system_prompts.ANALYST
|
||||
|
||||
|
||||
if self.decensor_enabled:
|
||||
|
||||
if "narrate" in kind:
|
||||
return system_prompts.NARRATOR
|
||||
if "story" in kind:
|
||||
return system_prompts.NARRATOR
|
||||
if "director" in kind:
|
||||
return system_prompts.DIRECTOR
|
||||
if "create" in kind:
|
||||
return system_prompts.CREATOR
|
||||
if "roleplay" in kind:
|
||||
return system_prompts.ROLEPLAY
|
||||
if "conversation" in kind:
|
||||
return system_prompts.ROLEPLAY
|
||||
if "editor" in kind:
|
||||
return system_prompts.EDITOR
|
||||
if "world_state" in kind:
|
||||
return system_prompts.WORLD_STATE
|
||||
if "analyze_freeform" in kind:
|
||||
return system_prompts.ANALYST_FREEFORM
|
||||
if "analyst" in kind:
|
||||
return system_prompts.ANALYST
|
||||
if "analyze" in kind:
|
||||
return system_prompts.ANALYST
|
||||
if "summarize" in kind:
|
||||
return system_prompts.SUMMARIZE
|
||||
if "visualize" in kind:
|
||||
return system_prompts.VISUALIZE
|
||||
|
||||
else:
|
||||
|
||||
if "narrate" in kind:
|
||||
return system_prompts.NARRATOR_NO_DECENSOR
|
||||
if "story" in kind:
|
||||
return system_prompts.NARRATOR_NO_DECENSOR
|
||||
if "director" in kind:
|
||||
return system_prompts.DIRECTOR_NO_DECENSOR
|
||||
if "create" in kind:
|
||||
return system_prompts.CREATOR_NO_DECENSOR
|
||||
if "roleplay" in kind:
|
||||
return system_prompts.ROLEPLAY_NO_DECENSOR
|
||||
if "conversation" in kind:
|
||||
return system_prompts.ROLEPLAY_NO_DECENSOR
|
||||
if "editor" in kind:
|
||||
return system_prompts.EDITOR_NO_DECENSOR
|
||||
if "world_state" in kind:
|
||||
return system_prompts.WORLD_STATE_NO_DECENSOR
|
||||
if "analyze_freeform" in kind:
|
||||
return system_prompts.ANALYST_FREEFORM_NO_DECENSOR
|
||||
if "analyst" in kind:
|
||||
return system_prompts.ANALYST_NO_DECENSOR
|
||||
if "analyze" in kind:
|
||||
return system_prompts.ANALYST_NO_DECENSOR
|
||||
if "summarize" in kind:
|
||||
return system_prompts.SUMMARIZE_NO_DECENSOR
|
||||
if "visualize" in kind:
|
||||
return system_prompts.VISUALIZE_NO_DECENSOR
|
||||
|
||||
return system_prompts.BASIC
|
||||
|
||||
|
||||
|
||||
def emit_status(self, processing: bool = None):
|
||||
|
||||
"""
|
||||
Sets and emits the client status.
|
||||
"""
|
||||
|
||||
|
||||
if processing is not None:
|
||||
self.processing = processing
|
||||
|
||||
@@ -171,29 +257,45 @@ class ClientBase:
|
||||
else:
|
||||
model_name = "No model loaded"
|
||||
status = "warning"
|
||||
|
||||
|
||||
status_change = status != self.current_status
|
||||
self.current_status = status
|
||||
|
||||
|
||||
prompt_template_example, prompt_template_file = self.prompt_template_example()
|
||||
|
||||
data = {
|
||||
"api_key": self.api_key,
|
||||
"prompt_template_example": prompt_template_example,
|
||||
"has_prompt_template": (
|
||||
prompt_template_file and prompt_template_file != "default.jinja2"
|
||||
),
|
||||
"template_file": prompt_template_file,
|
||||
"meta": self.Meta().model_dump(),
|
||||
"error_action": None,
|
||||
}
|
||||
|
||||
for field_name in getattr(self.Meta(), "extra_fields", {}).keys():
|
||||
data[field_name] = getattr(self, field_name, None)
|
||||
|
||||
emit(
|
||||
"client_status",
|
||||
message=self.client_type,
|
||||
id=self.name,
|
||||
details=model_name,
|
||||
status=status,
|
||||
data=data,
|
||||
)
|
||||
|
||||
if status_change:
|
||||
instance.emit_agent_status_by_client(self)
|
||||
|
||||
|
||||
|
||||
async def get_model_name(self):
|
||||
models = await self.client.models.list()
|
||||
try:
|
||||
return models.data[0].id
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
|
||||
async def status(self):
|
||||
"""
|
||||
Send a request to the API to retrieve the loaded AI model name.
|
||||
@@ -202,12 +304,12 @@ class ClientBase:
|
||||
"""
|
||||
if self.processing:
|
||||
return
|
||||
|
||||
|
||||
if not self.enabled:
|
||||
self.connected = False
|
||||
self.emit_status()
|
||||
return
|
||||
|
||||
|
||||
try:
|
||||
self.model_name = await self.get_model_name()
|
||||
except Exception as e:
|
||||
@@ -217,133 +319,331 @@ class ClientBase:
|
||||
self.toggle_disabled_if_remote()
|
||||
self.emit_status()
|
||||
return
|
||||
|
||||
|
||||
self.connected = True
|
||||
|
||||
|
||||
if not self.model_name or self.model_name == "None":
|
||||
self.log.warning("client model not loaded", client=self)
|
||||
self.emit_status()
|
||||
return
|
||||
|
||||
|
||||
self.emit_status()
|
||||
|
||||
|
||||
def generate_prompt_parameters(self, kind:str):
|
||||
|
||||
def generate_prompt_parameters(self, kind: str):
|
||||
parameters = {}
|
||||
self.tune_prompt_parameters(
|
||||
presets.configure(parameters, kind, self.max_token_length),
|
||||
kind
|
||||
presets.configure(parameters, kind, self.max_token_length), kind
|
||||
)
|
||||
return parameters
|
||||
|
||||
def tune_prompt_parameters(self, parameters:dict, kind:str):
|
||||
|
||||
def tune_prompt_parameters(self, parameters: dict, kind: str):
|
||||
parameters["stream"] = False
|
||||
if client_context_attribute("nuke_repetition") > 0.0 and self.jiggle_enabled_for(kind):
|
||||
self.jiggle_randomness(parameters, offset=client_context_attribute("nuke_repetition"))
|
||||
|
||||
if client_context_attribute(
|
||||
"nuke_repetition"
|
||||
) > 0.0 and self.jiggle_enabled_for(kind):
|
||||
self.jiggle_randomness(
|
||||
parameters, offset=client_context_attribute("nuke_repetition")
|
||||
)
|
||||
|
||||
fn_tune_kind = getattr(self, f"tune_prompt_parameters_{kind}", None)
|
||||
if fn_tune_kind:
|
||||
fn_tune_kind(parameters)
|
||||
|
||||
def tune_prompt_parameters_conversation(self, parameters:dict):
|
||||
|
||||
agent_context = active_agent.get()
|
||||
if agent_context.agent:
|
||||
agent_context.agent.inject_prompt_paramters(
|
||||
parameters, kind, agent_context.action
|
||||
)
|
||||
|
||||
def tune_prompt_parameters_conversation(self, parameters: dict):
|
||||
conversation_context = client_context_attribute("conversation")
|
||||
parameters["max_tokens"] = conversation_context.get("length", 96)
|
||||
|
||||
|
||||
dialog_stopping_strings = [
|
||||
f"{character}:" for character in conversation_context["other_characters"]
|
||||
]
|
||||
|
||||
|
||||
dialog_stopping_strings += [
|
||||
f"{character.upper()}\n"
|
||||
for character in conversation_context["other_characters"]
|
||||
]
|
||||
|
||||
if "extra_stopping_strings" in parameters:
|
||||
parameters["extra_stopping_strings"] += dialog_stopping_strings
|
||||
else:
|
||||
parameters["extra_stopping_strings"] = dialog_stopping_strings
|
||||
|
||||
|
||||
async def generate(self, prompt:str, parameters:dict, kind:str):
|
||||
|
||||
|
||||
async def generate(self, prompt: str, parameters: dict, kind: str):
|
||||
"""
|
||||
Generates text from the given prompt and parameters.
|
||||
"""
|
||||
|
||||
self.log.debug("generate", prompt=prompt[:128]+" ...", parameters=parameters)
|
||||
|
||||
|
||||
self.log.debug("generate", prompt=prompt[:128] + " ...", parameters=parameters)
|
||||
|
||||
try:
|
||||
response = await self.client.completions.create(prompt=prompt.strip(), **parameters)
|
||||
response = await self.client.completions.create(
|
||||
prompt=prompt.strip(" "), **parameters
|
||||
)
|
||||
return response.get("choices", [{}])[0].get("text", "")
|
||||
except PermissionDeniedError as e:
|
||||
self.log.error("generate error", e=e)
|
||||
emit("status", message="Client API: Permission Denied", status="error")
|
||||
return ""
|
||||
except Exception as e:
|
||||
self.log.error("generate error", e=e)
|
||||
emit(
|
||||
"status", message="Error during generation (check logs)", status="error"
|
||||
)
|
||||
return ""
|
||||
|
||||
|
||||
async def send_prompt(
|
||||
self, prompt: str, kind: str = "conversation", finalize: Callable = lambda x: x
|
||||
self,
|
||||
prompt: str,
|
||||
kind: str = "conversation",
|
||||
finalize: Callable = lambda x: x,
|
||||
retries: int = 2,
|
||||
) -> str:
|
||||
"""
|
||||
Send a prompt to the AI and return its response.
|
||||
:param prompt: The text prompt to send.
|
||||
:return: The AI's response text.
|
||||
"""
|
||||
|
||||
|
||||
try:
|
||||
self._returned_prompt_tokens = None
|
||||
self._returned_response_tokens = None
|
||||
|
||||
self.emit_status(processing=True)
|
||||
await self.status()
|
||||
|
||||
prompt_param = self.generate_prompt_parameters(kind)
|
||||
|
||||
finalized_prompt = self.prompt_template(self.get_system_message(kind), prompt).strip()
|
||||
finalized_prompt = self.prompt_template(
|
||||
self.get_system_message(kind), prompt
|
||||
).strip(" ")
|
||||
prompt_param = finalize(prompt_param)
|
||||
|
||||
token_length = self.count_tokens(finalized_prompt)
|
||||
|
||||
|
||||
time_start = time.time()
|
||||
extra_stopping_strings = prompt_param.pop("extra_stopping_strings", [])
|
||||
|
||||
self.log.debug("send_prompt", token_length=token_length, max_token_length=self.max_token_length, parameters=prompt_param)
|
||||
response = await self.generate(finalized_prompt, prompt_param, kind)
|
||||
|
||||
self.log.debug(
|
||||
"send_prompt",
|
||||
token_length=token_length,
|
||||
max_token_length=self.max_token_length,
|
||||
parameters=prompt_param,
|
||||
)
|
||||
response = await self.generate(
|
||||
self.repetition_adjustment(finalized_prompt), prompt_param, kind
|
||||
)
|
||||
|
||||
response, finalized_prompt = await self.auto_break_repetition(
|
||||
finalized_prompt, prompt_param, response, kind, retries
|
||||
)
|
||||
|
||||
time_end = time.time()
|
||||
|
||||
|
||||
# stopping strings sometimes get appended to the end of the response anyways
|
||||
# split the response by the first stopping string and take the first part
|
||||
|
||||
|
||||
|
||||
for stopping_string in STOPPING_STRINGS + extra_stopping_strings:
|
||||
if stopping_string in response:
|
||||
response = response.split(stopping_string)[0]
|
||||
break
|
||||
|
||||
emit("prompt_sent", data={
|
||||
"kind": kind,
|
||||
"prompt": finalized_prompt,
|
||||
"response": response,
|
||||
"prompt_tokens": token_length,
|
||||
"response_tokens": self.count_tokens(response),
|
||||
"time": time_end - time_start,
|
||||
})
|
||||
|
||||
|
||||
agent_context = active_agent.get()
|
||||
|
||||
emit(
|
||||
"prompt_sent",
|
||||
data=PromptData(
|
||||
kind=kind,
|
||||
prompt=finalized_prompt,
|
||||
response=response,
|
||||
prompt_tokens=self._returned_prompt_tokens or token_length,
|
||||
response_tokens=self._returned_response_tokens
|
||||
or self.count_tokens(response),
|
||||
agent_stack=agent_context.agent_stack if agent_context else [],
|
||||
client_name=self.name,
|
||||
client_type=self.client_type,
|
||||
time=time_end - time_start,
|
||||
generation_parameters=prompt_param,
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
return response
|
||||
finally:
|
||||
self.emit_status(processing=False)
|
||||
|
||||
def count_tokens(self, content:str):
|
||||
self._returned_prompt_tokens = None
|
||||
self._returned_response_tokens = None
|
||||
|
||||
async def auto_break_repetition(
|
||||
self,
|
||||
finalized_prompt: str,
|
||||
prompt_param: dict,
|
||||
response: str,
|
||||
kind: str,
|
||||
retries: int,
|
||||
pad_max_tokens: int = 32,
|
||||
) -> str:
|
||||
"""
|
||||
If repetition breaking is enabled, this will retry the prompt if its
|
||||
response is too similar to other messages in the prompt
|
||||
|
||||
This requires the agent to have the allow_repetition_break method
|
||||
and the jiggle_enabled_for method and the client to have the
|
||||
auto_break_repetition_enabled attribute set to True
|
||||
|
||||
Arguments:
|
||||
|
||||
- finalized_prompt: the prompt that was sent
|
||||
- prompt_param: the parameters that were used
|
||||
- response: the response that was received
|
||||
- kind: the kind of generation
|
||||
- retries: the number of retries left
|
||||
- pad_max_tokens: increase response max_tokens by this amount per iteration
|
||||
|
||||
Returns:
|
||||
|
||||
- the response
|
||||
"""
|
||||
|
||||
if not self.auto_break_repetition_enabled:
|
||||
return response, finalized_prompt
|
||||
|
||||
agent_context = active_agent.get()
|
||||
if self.jiggle_enabled_for(kind, auto=True):
|
||||
# check if the response is a repetition
|
||||
# using the default similarity threshold of 98, meaning it needs
|
||||
# to be really similar to be considered a repetition
|
||||
|
||||
is_repetition, similarity_score, matched_line = util.similarity_score(
|
||||
response, finalized_prompt.split("\n"), similarity_threshold=80
|
||||
)
|
||||
|
||||
if not is_repetition:
|
||||
# not a repetition, return the response
|
||||
|
||||
self.log.debug(
|
||||
"send_prompt no similarity", similarity_score=similarity_score
|
||||
)
|
||||
finalized_prompt = self.repetition_adjustment(
|
||||
finalized_prompt, is_repetitive=False
|
||||
)
|
||||
return response, finalized_prompt
|
||||
|
||||
while is_repetition and retries > 0:
|
||||
# it's a repetition, retry the prompt with adjusted parameters
|
||||
|
||||
self.log.warn(
|
||||
"send_prompt similarity retry",
|
||||
agent=agent_context.agent.agent_type,
|
||||
similarity_score=similarity_score,
|
||||
retries=retries,
|
||||
)
|
||||
|
||||
# first we apply the client's randomness jiggle which will adjust
|
||||
# parameters like temperature and repetition_penalty, depending
|
||||
# on the client
|
||||
#
|
||||
# this is a cumulative adjustment, so it will add to the previous
|
||||
# iteration's adjustment, this also means retries should be kept low
|
||||
# otherwise it will get out of hand and start generating nonsense
|
||||
|
||||
self.jiggle_randomness(prompt_param, offset=0.5)
|
||||
|
||||
# then we pad the max_tokens by the pad_max_tokens amount
|
||||
|
||||
prompt_param["max_tokens"] += pad_max_tokens
|
||||
|
||||
# send the prompt again
|
||||
# we use the repetition_adjustment method to further encourage
|
||||
# the AI to break the repetition on its own as well.
|
||||
|
||||
finalized_prompt = self.repetition_adjustment(
|
||||
finalized_prompt, is_repetitive=True
|
||||
)
|
||||
|
||||
response = retried_response = await self.generate(
|
||||
finalized_prompt, prompt_param, kind
|
||||
)
|
||||
|
||||
self.log.debug(
|
||||
"send_prompt dedupe sentences",
|
||||
response=response,
|
||||
matched_line=matched_line,
|
||||
)
|
||||
|
||||
# a lot of the times the response will now contain the repetition + something new
|
||||
# so we dedupe the response to remove the repetition on sentences level
|
||||
|
||||
response = util.dedupe_sentences(
|
||||
response, matched_line, similarity_threshold=85, debug=True
|
||||
)
|
||||
self.log.debug(
|
||||
"send_prompt dedupe sentences (after)", response=response
|
||||
)
|
||||
|
||||
# deduping may have removed the entire response, so we check for that
|
||||
|
||||
if not util.strip_partial_sentences(response).strip():
|
||||
# if the response is empty, we set the response to the original
|
||||
# and try again next loop
|
||||
|
||||
response = retried_response
|
||||
|
||||
# check if the response is a repetition again
|
||||
|
||||
is_repetition, similarity_score, matched_line = util.similarity_score(
|
||||
response, finalized_prompt.split("\n"), similarity_threshold=80
|
||||
)
|
||||
retries -= 1
|
||||
|
||||
return response, finalized_prompt
|
||||
|
||||
def count_tokens(self, content: str):
|
||||
return util.count_tokens(content)
|
||||
|
||||
def jiggle_randomness(self, prompt_config:dict, offset:float=0.3) -> dict:
|
||||
def jiggle_randomness(self, prompt_config: dict, offset: float = 0.3) -> dict:
|
||||
"""
|
||||
adjusts temperature and repetition_penalty
|
||||
by random values using the base value as a center
|
||||
"""
|
||||
|
||||
|
||||
temp = prompt_config["temperature"]
|
||||
min_offset = offset * 0.3
|
||||
prompt_config["temperature"] = random.uniform(temp + min_offset, temp + offset)
|
||||
|
||||
def jiggle_enabled_for(self, kind:str):
|
||||
|
||||
if kind in ["conversation", "story"]:
|
||||
return True
|
||||
|
||||
if kind.startswith("narrate"):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def jiggle_enabled_for(self, kind: str, auto: bool = False) -> bool:
|
||||
agent_context = active_agent.get()
|
||||
agent = agent_context.agent
|
||||
|
||||
if not agent:
|
||||
return False
|
||||
|
||||
return agent.allow_repetition_break(kind, agent_context.action, auto=auto)
|
||||
|
||||
def repetition_adjustment(self, prompt: str, is_repetitive: bool = False):
|
||||
"""
|
||||
Breaks the prompt into lines and checkse each line for a match with
|
||||
[$REPETITION|{repetition_adjustment}].
|
||||
|
||||
On match and if is_repetitive is True, the line is removed from the prompt and
|
||||
replaced with the repetition_adjustment.
|
||||
|
||||
On match and if is_repetitive is False, the line is removed from the prompt.
|
||||
"""
|
||||
|
||||
lines = prompt.split("\n")
|
||||
new_lines = []
|
||||
|
||||
for line in lines:
|
||||
if line.startswith("[$REPETITION|"):
|
||||
if is_repetitive:
|
||||
new_lines.append(line.split("|")[1][:-1])
|
||||
else:
|
||||
new_lines.append("")
|
||||
else:
|
||||
new_lines.append(line)
|
||||
|
||||
return "\n".join(new_lines)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import pydantic
|
||||
from enum import Enum
|
||||
|
||||
import pydantic
|
||||
|
||||
__all__ = [
|
||||
"ClientType",
|
||||
"ClientBootstrap",
|
||||
@@ -10,8 +11,10 @@ __all__ = [
|
||||
|
||||
LISTS = {}
|
||||
|
||||
|
||||
class ClientType(str, Enum):
|
||||
"""Client type enum."""
|
||||
|
||||
textgen = "textgenwebui"
|
||||
automatic1111 = "automatic1111"
|
||||
|
||||
@@ -20,43 +23,42 @@ class ClientBootstrap(pydantic.BaseModel):
|
||||
"""Client bootstrap model."""
|
||||
|
||||
# client type, currently supports "textgen" and "automatic1111"
|
||||
|
||||
|
||||
client_type: ClientType
|
||||
|
||||
|
||||
# unique client identifier
|
||||
|
||||
|
||||
uid: str
|
||||
|
||||
|
||||
# connection name
|
||||
|
||||
|
||||
name: str
|
||||
|
||||
|
||||
# connection information for the client
|
||||
# REST api url
|
||||
|
||||
|
||||
api_url: str
|
||||
|
||||
|
||||
# service name (for example runpod)
|
||||
|
||||
|
||||
service_name: str
|
||||
|
||||
|
||||
|
||||
class register_list:
|
||||
|
||||
def __init__(self, service_name:str):
|
||||
def __init__(self, service_name: str):
|
||||
self.service_name = service_name
|
||||
|
||||
|
||||
def __call__(self, func):
|
||||
LISTS[self.service_name] = func
|
||||
return func
|
||||
|
||||
|
||||
def list_all(exclude_urls: list[str] = list()):
|
||||
|
||||
async def list_all(exclude_urls: list[str] = list()):
|
||||
"""
|
||||
Return a list of client bootstrap objects.
|
||||
"""
|
||||
|
||||
|
||||
for service_name, func in LISTS.items():
|
||||
for item in func():
|
||||
async for item in func():
|
||||
if item.api_url not in exclude_urls:
|
||||
yield item.dict()
|
||||
yield item.dict()
|
||||
|
||||
@@ -3,19 +3,20 @@ Context managers for various client-side operations.
|
||||
"""
|
||||
|
||||
from contextvars import ContextVar
|
||||
from pydantic import BaseModel, Field
|
||||
from copy import deepcopy
|
||||
|
||||
import structlog
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
__all__ = [
|
||||
'context_data',
|
||||
'client_context_attribute',
|
||||
'ContextModel',
|
||||
"context_data",
|
||||
"client_context_attribute",
|
||||
"ContextModel",
|
||||
]
|
||||
|
||||
log = structlog.get_logger()
|
||||
|
||||
|
||||
def model_to_dict_without_defaults(model_instance):
|
||||
model_dict = model_instance.dict()
|
||||
for field_name, field in model_instance.__class__.__fields__.items():
|
||||
@@ -23,20 +24,25 @@ def model_to_dict_without_defaults(model_instance):
|
||||
del model_dict[field_name]
|
||||
return model_dict
|
||||
|
||||
|
||||
class ConversationContext(BaseModel):
|
||||
talking_character: str = None
|
||||
other_characters: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ContextModel(BaseModel):
|
||||
"""
|
||||
Pydantic model for the context data.
|
||||
"""
|
||||
|
||||
nuke_repetition: float = Field(0.0, ge=0.0, le=3.0)
|
||||
conversation: ConversationContext = Field(default_factory=ConversationContext)
|
||||
length: int = 96
|
||||
|
||||
|
||||
# Define the context variable as an empty dictionary
|
||||
context_data = ContextVar('context_data', default=ContextModel().model_dump())
|
||||
context_data = ContextVar("context_data", default=ContextModel().model_dump())
|
||||
|
||||
|
||||
def client_context_attribute(name, default=None):
|
||||
"""
|
||||
@@ -47,6 +53,7 @@ def client_context_attribute(name, default=None):
|
||||
# Return the value of the key if it exists, otherwise return the default value
|
||||
return data.get(name, default)
|
||||
|
||||
|
||||
def set_client_context_attribute(name, value):
|
||||
"""
|
||||
Set the value of the context variable `context_data` for the given key.
|
||||
@@ -55,7 +62,8 @@ def set_client_context_attribute(name, value):
|
||||
data = context_data.get()
|
||||
# Set the value of the key
|
||||
data[name] = value
|
||||
|
||||
|
||||
|
||||
def set_conversation_context_attribute(name, value):
|
||||
"""
|
||||
Set the value of the context variable `context_data.conversation` for the given key.
|
||||
@@ -65,6 +73,7 @@ def set_conversation_context_attribute(name, value):
|
||||
# Set the value of the key
|
||||
data["conversation"][name] = value
|
||||
|
||||
|
||||
class ClientContext:
|
||||
"""
|
||||
A context manager to set values to the context variable `context_data`.
|
||||
@@ -82,10 +91,10 @@ class ClientContext:
|
||||
Set the key-value pairs to the context variable `context_data` when entering the context.
|
||||
"""
|
||||
# Get the current context data
|
||||
|
||||
|
||||
data = deepcopy(context_data.get()) if context_data.get() else {}
|
||||
data.update(self.values)
|
||||
|
||||
|
||||
# Update the context data
|
||||
self.token = context_data.set(data)
|
||||
|
||||
@@ -93,5 +102,5 @@ class ClientContext:
|
||||
"""
|
||||
Reset the context variable `context_data` to its previous values when exiting the context.
|
||||
"""
|
||||
|
||||
|
||||
context_data.reset(self.token)
|
||||
|
||||
34
src/talemate/client/custom/__init__.py
Normal file
@@ -0,0 +1,34 @@
|
||||
import importlib
|
||||
import os
|
||||
|
||||
import structlog
|
||||
|
||||
log = structlog.get_logger("talemate.client.custom")
|
||||
|
||||
# import every submodule in this directory
|
||||
#
|
||||
# each directory in this directory is a submodule
|
||||
|
||||
# get the current directory
|
||||
current_directory = os.path.dirname(__file__)
|
||||
|
||||
# get all subdirectories
|
||||
subdirectories = [
|
||||
os.path.join(current_directory, name)
|
||||
for name in os.listdir(current_directory)
|
||||
if os.path.isdir(os.path.join(current_directory, name))
|
||||
]
|
||||
|
||||
# import every submodule
|
||||
|
||||
for subdirectory in subdirectories:
|
||||
# get the name of the submodule
|
||||
submodule_name = os.path.basename(subdirectory)
|
||||
|
||||
if submodule_name.startswith("__"):
|
||||
continue
|
||||
|
||||
log.info("activating custom client", module=submodule_name)
|
||||
|
||||
# import the submodule
|
||||
importlib.import_module(f".{submodule_name}", __package__)
|
||||
@@ -0,0 +1,5 @@
|
||||
Each client should be in its own subdirectory.
|
||||
|
||||
The subdirectory itself must be a valid python module.
|
||||
|
||||
Check out docs/dev/client/example/test for a very simplistic custom client example.
|
||||
@@ -1,16 +1,16 @@
|
||||
import asyncio
|
||||
import random
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable, Union
|
||||
|
||||
import requests
|
||||
|
||||
import talemate.client.system_prompts as system_prompts
|
||||
import talemate.util as util
|
||||
from talemate.client.registry import register
|
||||
import talemate.client.system_prompts as system_prompts
|
||||
from talemate.client.textgenwebui import RESTTaleMateClient
|
||||
from talemate.emit import Emission, emit
|
||||
|
||||
# NOT IMPLEMENTED AT THIS POINT
|
||||
# NOT IMPLEMENTED AT THIS POINT
|
||||
|
||||
@@ -1,56 +1,63 @@
|
||||
import pydantic
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from talemate.client.base import ClientBase
|
||||
from talemate.client.registry import register
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
class Defaults(pydantic.BaseModel):
|
||||
api_url: str = "http://localhost:1234"
|
||||
max_token_length: int = 4096
|
||||
|
||||
|
||||
@register()
|
||||
class LMStudioClient(ClientBase):
|
||||
|
||||
client_type = "lmstudio"
|
||||
conversation_retries = 5
|
||||
|
||||
def set_client(self):
|
||||
self.client = AsyncOpenAI(base_url=self.api_url+"/v1", api_key="sk-1111")
|
||||
|
||||
def tune_prompt_parameters(self, parameters:dict, kind:str):
|
||||
class Meta(ClientBase.Meta):
|
||||
name_prefix: str = "LMStudio"
|
||||
title: str = "LMStudio"
|
||||
defaults: Defaults = Defaults()
|
||||
|
||||
def set_client(self, **kwargs):
|
||||
self.client = AsyncOpenAI(base_url=self.api_url + "/v1", api_key="sk-1111")
|
||||
|
||||
def tune_prompt_parameters(self, parameters: dict, kind: str):
|
||||
super().tune_prompt_parameters(parameters, kind)
|
||||
|
||||
|
||||
keys = list(parameters.keys())
|
||||
|
||||
|
||||
valid_keys = ["temperature", "top_p"]
|
||||
|
||||
|
||||
for key in keys:
|
||||
if key not in valid_keys:
|
||||
del parameters[key]
|
||||
|
||||
|
||||
|
||||
async def get_model_name(self):
|
||||
model_name = await super().get_model_name()
|
||||
|
||||
# model name comes back as a file path, so we need to extract the model name
|
||||
# the path could be windows or linux so it needs to handle both backslash and forward slash
|
||||
|
||||
|
||||
if model_name:
|
||||
model_name = model_name.replace("\\", "/").split("/")[-1]
|
||||
|
||||
return model_name
|
||||
|
||||
async def generate(self, prompt:str, parameters:dict, kind:str):
|
||||
|
||||
async def generate(self, prompt: str, parameters: dict, kind: str):
|
||||
"""
|
||||
Generates text from the given prompt and parameters.
|
||||
"""
|
||||
human_message = {'role': 'user', 'content': prompt.strip()}
|
||||
|
||||
self.log.debug("generate", prompt=prompt[:128]+" ...", parameters=parameters)
|
||||
|
||||
human_message = {"role": "user", "content": prompt.strip()}
|
||||
|
||||
self.log.debug("generate", prompt=prompt[:128] + " ...", parameters=parameters)
|
||||
|
||||
try:
|
||||
response = await self.client.chat.completions.create(
|
||||
model=self.model_name, messages=[human_message], **parameters
|
||||
)
|
||||
|
||||
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
self.log.error("generate error", e=e)
|
||||
return ""
|
||||
return ""
|
||||
|
||||
232
src/talemate/client/mistral.py
Normal file
@@ -0,0 +1,232 @@
|
||||
import json
|
||||
|
||||
import pydantic
|
||||
import structlog
|
||||
import tiktoken
|
||||
from openai import AsyncOpenAI, PermissionDeniedError
|
||||
|
||||
from talemate.client.base import ClientBase, ErrorAction
|
||||
from talemate.client.registry import register
|
||||
from talemate.config import load_config
|
||||
from talemate.emit import emit
|
||||
from talemate.emit.signals import handlers
|
||||
|
||||
__all__ = [
|
||||
"MistralAIClient",
|
||||
]
|
||||
log = structlog.get_logger("talemate")
|
||||
|
||||
# Edit this to add new models / remove old models
|
||||
SUPPORTED_MODELS = [
|
||||
"open-mistral-7b",
|
||||
"open-mixtral-8x7b",
|
||||
"mistral-small-latest",
|
||||
"mistral-medium-latest",
|
||||
"mistral-large-latest",
|
||||
]
|
||||
|
||||
|
||||
class Defaults(pydantic.BaseModel):
|
||||
max_token_length: int = 16384
|
||||
model: str = "open-mixtral-8x7b"
|
||||
|
||||
|
||||
@register()
|
||||
class MistralAIClient(ClientBase):
|
||||
"""
|
||||
OpenAI client for generating text.
|
||||
"""
|
||||
|
||||
client_type = "mistral"
|
||||
conversation_retries = 0
|
||||
auto_break_repetition_enabled = False
|
||||
# TODO: make this configurable?
|
||||
decensor_enabled = False
|
||||
|
||||
class Meta(ClientBase.Meta):
|
||||
name_prefix: str = "MistralAI"
|
||||
title: str = "MistralAI"
|
||||
manual_model: bool = True
|
||||
manual_model_choices: list[str] = SUPPORTED_MODELS
|
||||
requires_prompt_template: bool = False
|
||||
defaults: Defaults = Defaults()
|
||||
|
||||
def __init__(self, model="open-mixtral-8x7b", **kwargs):
|
||||
self.model_name = model
|
||||
self.api_key_status = None
|
||||
self.config = load_config()
|
||||
super().__init__(**kwargs)
|
||||
|
||||
handlers["config_saved"].connect(self.on_config_saved)
|
||||
|
||||
@property
|
||||
def mistralai_api_key(self):
|
||||
return self.config.get("mistralai", {}).get("api_key")
|
||||
|
||||
def emit_status(self, processing: bool = None):
|
||||
error_action = None
|
||||
if processing is not None:
|
||||
self.processing = processing
|
||||
|
||||
if self.mistralai_api_key:
|
||||
status = "busy" if self.processing else "idle"
|
||||
model_name = self.model_name
|
||||
else:
|
||||
status = "error"
|
||||
model_name = "No API key set"
|
||||
error_action = ErrorAction(
|
||||
title="Set API Key",
|
||||
action_name="openAppConfig",
|
||||
icon="mdi-key-variant",
|
||||
arguments=[
|
||||
"application",
|
||||
"mistralai_api",
|
||||
],
|
||||
)
|
||||
|
||||
if not self.model_name:
|
||||
status = "error"
|
||||
model_name = "No model loaded"
|
||||
|
||||
self.current_status = status
|
||||
|
||||
emit(
|
||||
"client_status",
|
||||
message=self.client_type,
|
||||
id=self.name,
|
||||
details=model_name,
|
||||
status=status,
|
||||
data={
|
||||
"error_action": error_action.model_dump() if error_action else None,
|
||||
"meta": self.Meta().model_dump(),
|
||||
},
|
||||
)
|
||||
|
||||
def set_client(self, max_token_length: int = None):
|
||||
if not self.mistralai_api_key:
|
||||
self.client = AsyncOpenAI(api_key="sk-1111")
|
||||
log.error("No mistral.ai API key set")
|
||||
if self.api_key_status:
|
||||
self.api_key_status = False
|
||||
emit("request_client_status")
|
||||
emit("request_agent_status")
|
||||
return
|
||||
|
||||
if not self.model_name:
|
||||
self.model_name = "open-mixtral-8x7b"
|
||||
|
||||
if max_token_length and not isinstance(max_token_length, int):
|
||||
max_token_length = int(max_token_length)
|
||||
|
||||
model = self.model_name
|
||||
|
||||
self.client = AsyncOpenAI(
|
||||
api_key=self.mistralai_api_key, base_url="https://api.mistral.ai/v1/"
|
||||
)
|
||||
self.max_token_length = max_token_length or 16384
|
||||
|
||||
if not self.api_key_status:
|
||||
if self.api_key_status is False:
|
||||
emit("request_client_status")
|
||||
emit("request_agent_status")
|
||||
self.api_key_status = True
|
||||
|
||||
log.info(
|
||||
"mistral.ai set client",
|
||||
max_token_length=self.max_token_length,
|
||||
provided_max_token_length=max_token_length,
|
||||
model=model,
|
||||
)
|
||||
|
||||
def reconfigure(self, **kwargs):
|
||||
if kwargs.get("model"):
|
||||
self.model_name = kwargs["model"]
|
||||
self.set_client(kwargs.get("max_token_length"))
|
||||
|
||||
def on_config_saved(self, event):
|
||||
config = event.data
|
||||
self.config = config
|
||||
self.set_client(max_token_length=self.max_token_length)
|
||||
|
||||
def response_tokens(self, response: str):
|
||||
return response.usage.completion_tokens
|
||||
|
||||
def prompt_tokens(self, response: str):
|
||||
return response.usage.prompt_tokens
|
||||
|
||||
async def status(self):
|
||||
self.emit_status()
|
||||
|
||||
def prompt_template(self, system_message: str, prompt: str):
|
||||
if "<|BOT|>" in prompt:
|
||||
_, right = prompt.split("<|BOT|>", 1)
|
||||
if right:
|
||||
prompt = prompt.replace("<|BOT|>", "\nStart your response with: ")
|
||||
else:
|
||||
prompt = prompt.replace("<|BOT|>", "")
|
||||
|
||||
return prompt
|
||||
|
||||
def tune_prompt_parameters(self, parameters: dict, kind: str):
|
||||
super().tune_prompt_parameters(parameters, kind)
|
||||
keys = list(parameters.keys())
|
||||
valid_keys = ["temperature", "top_p", "max_tokens"]
|
||||
for key in keys:
|
||||
if key not in valid_keys:
|
||||
del parameters[key]
|
||||
|
||||
async def generate(self, prompt: str, parameters: dict, kind: str):
|
||||
"""
|
||||
Generates text from the given prompt and parameters.
|
||||
"""
|
||||
|
||||
if not self.mistralai_api_key:
|
||||
raise Exception("No mistral.ai API key set")
|
||||
|
||||
right = None
|
||||
expected_response = None
|
||||
try:
|
||||
_, right = prompt.split("\nStart your response with: ")
|
||||
expected_response = right.strip()
|
||||
except (IndexError, ValueError):
|
||||
pass
|
||||
|
||||
human_message = {"role": "user", "content": prompt.strip()}
|
||||
system_message = {"role": "system", "content": self.get_system_message(kind)}
|
||||
|
||||
self.log.debug(
|
||||
"generate",
|
||||
prompt=prompt[:128] + " ...",
|
||||
parameters=parameters,
|
||||
system_message=system_message,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=[system_message, human_message],
|
||||
**parameters,
|
||||
)
|
||||
|
||||
self._returned_prompt_tokens = self.prompt_tokens(response)
|
||||
self._returned_response_tokens = self.response_tokens(response)
|
||||
|
||||
response = response.choices[0].message.content
|
||||
|
||||
# older models don't support json_object response coersion
|
||||
# and often like to return the response wrapped in ```json
|
||||
# so we strip that out if the expected response is a json object
|
||||
if expected_response and expected_response.startswith("{"):
|
||||
if response.startswith("```json") and response.endswith("```"):
|
||||
response = response[7:-3].strip()
|
||||
|
||||
if right and response.startswith(right):
|
||||
response = response[len(right) :].strip()
|
||||
|
||||
return response
|
||||
except PermissionDeniedError as e:
|
||||
self.log.error("generate error", e=e)
|
||||
emit("status", message="mistral.ai API: Permission Denied", status="error")
|
||||
return ""
|
||||
except Exception as e:
|
||||
raise
|
||||