Spaces:
Running
Running
File size: 5,129 Bytes
3a09141 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
import json
from typing import Literal
from pydantic import BaseModel
import litellm
from litellm.types.utils import ModelResponse
SAFETY_SETTINGS = [
{
"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE",
},
]
class BloggerAgent:
instructions = [
{
"role": "user",
"content": "与えられる情報について、重要なポイントを踏まえて平易な言葉で解説・紹介する記事を書いてください",
},
]
model: str = "gemini/gemini-2.5-flash-preview-05-20"
temperature: float = 1.0
max_tokens: int = 4096
thinking_budget: int = 1024
api_key: str
def __init__(self, api_key: str):
self.api_key = api_key
async def task(self, information: str) -> str:
messages = self.instructions.copy()
messages.append({"role": "user", "content": information})
res = await litellm.acompletion(
api_key=self.api_key,
model=self.model,
messages=messages,
temperature=self.temperature,
max_completion_tokens=self.max_tokens,
thinking={"type": "enabled", "budget_tokens": self.thinking_budget},
safety_settings=SAFETY_SETTINGS,
)
assert isinstance(res, ModelResponse)
blog = res.choices[0].message.content
assert isinstance(blog, str)
return blog
class WriterAgent:
instructions = [
{
"role": "user",
"content": """与えられる情報ソースとその解説記事をもとに、コンテンツを紹介する Podcast の会話を作成してください。
Podcast では、二人の人物が交互に会話をします。
# 登場人物
- スピーカー: コンテンツ紹介をリードする人で、主にこの人物が解説を行う
- サポーター: スピーカーの説明を聞き、うなづいたり、さらに質問を投げかけることで、理解を助ける。
# 構成
1. イントロ: まず、スピーカーとサポーターが何について話すのか、挨拶を交えながら会話します。自己紹介は省略する。
2. 解説: 前提知識の確認をしながら、内容を解説していきます
3. アウトロ: 今後の展望を交えながら締めくくります
---
このような内容になるような Podcast の脚本を作成してください。
""".strip(),
},
]
model: str = "gemini/gemini-2.5-flash-preview-05-20"
temperature: float = 1.0
max_tokens: int = 4096
thinking_budget: int = 1024
api_key: str
def __init__(self, api_key: str):
self.api_key = api_key
async def task(self, information: str, blog: str) -> str:
messages = self.instructions.copy()
messages.append(
{"role": "user", "content": f"# 情報\n{information}\n\n# 解説\n{blog}"}
)
res = await litellm.acompletion(
api_key=self.api_key,
model=self.model,
messages=messages,
temperature=self.temperature,
max_completion_tokens=self.max_tokens,
thinking={"type": "enabled", "budget_tokens": self.thinking_budget},
safety_settings=SAFETY_SETTINGS,
)
assert isinstance(res, ModelResponse)
dialogue = res.choices[0].message.content
assert isinstance(dialogue, str)
return dialogue
class Dialogue(BaseModel):
role: Literal["speaker", "supporter"]
content: str
class Conversation(BaseModel):
conversation: list[Dialogue]
class StructureAgent:
instructions = [
{
"role": "user",
"content": """この会話を指定されたスキーマに従った形に変換してください。スピーカーの role は `speaker`、サポーターは `supporter` です。""".strip(),
},
]
model: str = "gemini/gemini-2.5-flash-preview-05-20"
temperature: float = 0.1
max_tokens: int = 12_288
thinking_budget: int = 0
api_key: str
def __init__(self, api_key: str):
self.api_key = api_key
async def task(self, dialogue: str) -> Conversation:
messages = self.instructions.copy()
messages.append({"role": "user", "content": dialogue})
res = await litellm.acompletion(
api_key=self.api_key,
model=self.model,
messages=messages,
temperature=self.temperature,
max_completion_tokens=self.max_tokens,
thinking={"type": "disabled"},
response_format=Conversation,
safety_settings=SAFETY_SETTINGS,
)
conversation = Conversation.model_validate(
json.loads(res.choices[0].message.content)
)
return conversation
|