70 lines
2.1 KiB
Plaintext
70 lines
2.1 KiB
Plaintext
"""
|
|
title: Better R1
|
|
author: Aaron Feng
|
|
author_url: https://github.com/AaronFeng753
|
|
version: 0.2
|
|
"""
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing import Optional
|
|
import re
|
|
import time
|
|
|
|
|
|
class Filter:
|
|
start_time = None # class variable to track time
|
|
|
|
class Valves(BaseModel):
|
|
priority: int = Field(default=0, description="")
|
|
max_turns: int = Field(default=9999, description="")
|
|
|
|
class UserValves(BaseModel):
|
|
max_turns: int = Field(default=9999, description="")
|
|
|
|
def __init__(self):
|
|
self.valves = self.Valves()
|
|
|
|
def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
|
|
messages = body.get("messages", [])
|
|
modified_messages = []
|
|
|
|
for message in messages:
|
|
if message.get("role") == "assistant":
|
|
message_content = message.get("content", "")
|
|
pattern = r"<details[^>]*>.*?</details>"
|
|
modified_content = re.sub(pattern, "", message_content, flags=re.DOTALL)
|
|
modified_message = {"role": "assistant", "content": modified_content}
|
|
else:
|
|
modified_message = message
|
|
modified_messages.append(modified_message)
|
|
|
|
body["messages"] = modified_messages
|
|
|
|
Filter.start_time = time.time() # set initial time when instance created
|
|
|
|
return body
|
|
|
|
async def outlet(
|
|
self, body: dict, __user__: Optional[dict] = None, __event_emitter__=None
|
|
) -> dict:
|
|
elapsed_time = time.time() - Filter.start_time
|
|
|
|
messages = body.get("messages", [])
|
|
if not messages:
|
|
return body
|
|
|
|
model_response = messages[-1].get("content", "")
|
|
|
|
model_response = model_response.replace(
|
|
"<think>",
|
|
f"<details>\n<summary>Thought for {int(elapsed_time)} seconds 💭</summary>",
|
|
)
|
|
model_response = model_response.replace("</think>", "\n\n---\n\n</details>")
|
|
|
|
messages[-1]["content"] = model_response
|
|
body["messages"] = messages
|
|
|
|
Filter.start_time = time.time() # set initial time when instance created
|
|
|
|
return body
|