Refactor the code directory structure to better suit the current complexity
This commit is contained in:
@@ -6,9 +6,7 @@ import sys
|
||||
import runpy
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from aios_kernel import AIAgent,AIAgentTemplete,AIStorage,Environment
|
||||
from aios_kernel.agent_base import BaseAIAgent
|
||||
from package_manager import PackageEnv,PackageEnvManager,PackageMediaInfo,PackageInstallTask
|
||||
from aios import AIAgent,AIAgentTemplete,AIStorage,Environment,BaseAIAgent,PackageEnv,PackageEnvManager,PackageMediaInfo,PackageInstallTask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
import asyncio
|
||||
import aiosmtplib
|
||||
import aioimaplib
|
||||
import email
|
||||
from email.header import decode_header
|
||||
import mailparser
|
||||
import logging
|
||||
import time
|
||||
import datetime
|
||||
|
||||
from aios import AgentTunnel,AgentMsg,ContactManager
|
||||
|
||||
from email.message import EmailMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class EmailTunnel(AgentTunnel):
|
||||
@classmethod
|
||||
def register_to_loader(cls):
|
||||
async def load_email_tunnel(config:dict) -> AgentTunnel:
|
||||
result_tunnel = EmailTunnel()
|
||||
if await result_tunnel.load_from_config(config):
|
||||
return result_tunnel
|
||||
else:
|
||||
return None
|
||||
|
||||
AgentTunnel.register_loader("EmailTunnel",load_email_tunnel)
|
||||
|
||||
async def load_from_config(self,config:dict)->bool:
|
||||
self.target_id = config["target"]
|
||||
self.tunnel_id = config["tunnel_id"]
|
||||
|
||||
self.type = "EmailTunnel"
|
||||
self.email = config["email"]
|
||||
self.imap_server = config["imap"]
|
||||
s = self.imap_server.split(":")
|
||||
if len(s) == 2:
|
||||
self.imap_server = s[0]
|
||||
self.imap_port = int(s[1])
|
||||
|
||||
self.smtp_server = config["smtp"]
|
||||
s = self.smtp_server.split(":")
|
||||
if len(s) == 2:
|
||||
self.smtp_server = s[0]
|
||||
self.smtp_port = int(s[1])
|
||||
|
||||
self.login_user = config["user"]
|
||||
self.login_password = config["password"]
|
||||
if config.get("folder") is not None:
|
||||
self.folder = config["folder"]
|
||||
if config.get("interval") is not None:
|
||||
self.check_interval = config["interval"]
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.is_start = False
|
||||
self.read_email = {}
|
||||
self.folder = "INBOX"
|
||||
self.check_interval = 60
|
||||
|
||||
async def on_new_email(self,mail:mailparser.MailParser) -> None:
|
||||
remote_email_addr = mail.from_[0][1]
|
||||
remote_user_name = remote_email_addr.split("@")[0]
|
||||
agent_msg = self.conver_mail_to_agent_msg(mail)
|
||||
agent_msg.sender = remote_user_name
|
||||
agent_msg.target = self.target_id
|
||||
self.ai_bus.register_message_handler(remote_user_name, self._process_message)
|
||||
|
||||
resp_msg = await self.ai_bus.send_message(agent_msg)
|
||||
if resp_msg is None:
|
||||
await self.reply_email(remote_email_addr,"Sorry, I can't understand your message","")
|
||||
else:
|
||||
if resp_msg.body_mime is None:
|
||||
await self.reply_email(remote_email_addr,"result",resp_msg.body)
|
||||
|
||||
async def reply_email(self,target_email:str,title:str,msg:str) -> None:
|
||||
email_msg = EmailMessage()
|
||||
email_msg['Subject'] = f"Reply: {title}"
|
||||
email_msg['From'] = self.email
|
||||
email_msg['To'] = target_email
|
||||
email_msg.set_content(msg)
|
||||
|
||||
await aiosmtplib.send(
|
||||
email_msg,
|
||||
hostname = self.smtp_server,
|
||||
port=self.smtp_port,
|
||||
username=self.login_user,
|
||||
password=self.login_password,
|
||||
)
|
||||
|
||||
async def post_message(self, msg: AgentMsg) -> None:
|
||||
cm = ContactManager.get_instance()
|
||||
contact = cm.find_contact_by_name(msg.target)
|
||||
if contact is None:
|
||||
logger.error(f"can't find contact {msg.target} , post message through email_tunnel failed!")
|
||||
return
|
||||
|
||||
target_email = contact.email
|
||||
if target_email is None:
|
||||
logger.error(f"contact {msg.target} has no email, post message through email_tunnel failed!")
|
||||
return
|
||||
|
||||
email_msg = EmailMessage()
|
||||
email_msg['Subject'] = f"{msg.topic},From AIAgent {msg.sender}"
|
||||
email_msg['From'] = self.email
|
||||
email_msg['To'] = target_email
|
||||
email_msg.set_content(msg)
|
||||
|
||||
await aiosmtplib.send(
|
||||
email_msg,
|
||||
hostname = self.smtp_server,
|
||||
port=self.smtp_port,
|
||||
username=self.login_user,
|
||||
password=self.login_password,
|
||||
)
|
||||
|
||||
|
||||
def conver_mail_to_agent_msg(self,mail:mailparser.MailParser) -> AgentMsg:
|
||||
msg = AgentMsg()
|
||||
msg.set("",self.target_id,mail.text_plain[0])
|
||||
msg.topic = "email"
|
||||
return msg
|
||||
|
||||
async def check_email(self) -> None:
|
||||
self.last_check_num = 0
|
||||
self.last_check_time = datetime.datetime.now()
|
||||
while True:
|
||||
if self.is_start == False:
|
||||
return
|
||||
|
||||
await asyncio.sleep(self.check_interval)
|
||||
imap_client = aioimaplib.IMAP4_SSL(host=self.imap_server,port=self.imap_port)
|
||||
await imap_client.wait_hello_from_server()
|
||||
await imap_client.login(self.login_user, self.login_password)
|
||||
|
||||
date_since = self.last_check_time.strftime("%d-%b-%Y")
|
||||
|
||||
await imap_client.select(self.folder)
|
||||
status, messages = await imap_client.search('UNSEEN',charset='US-ASCII')
|
||||
self.last_check_time = datetime.datetime.now()
|
||||
if status == "OK":
|
||||
message_numbers = messages[0].split()
|
||||
for num in message_numbers:
|
||||
num = int(num)
|
||||
if self.read_email.get(num) is not None:
|
||||
continue
|
||||
|
||||
status, email_data = await imap_client.fetch(str(num), "(RFC822)")
|
||||
if status == "OK":
|
||||
#r = email.message_from_bytes(email_data[1])
|
||||
mail = mailparser.parse_from_bytes(email_data[1])
|
||||
self.read_email[num] = mail
|
||||
await self.on_new_email(mail)
|
||||
|
||||
await imap_client.logout()
|
||||
|
||||
async def start(self) -> bool:
|
||||
if self.is_start:
|
||||
logger.warning(f"tunnel {self.tunnel_id} is already started")
|
||||
return False
|
||||
self.is_start = True
|
||||
|
||||
asyncio.create_task(self.check_email())
|
||||
return True
|
||||
|
||||
async def close(self) -> None:
|
||||
self.is_start = False
|
||||
|
||||
async def _process_message(self, msg: AgentMsg) -> None:
|
||||
logger.warn(f"process message {msg.msg_id} from {msg.sender} to {msg.target}")
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .google_text_to_speech_node import *
|
||||
@@ -0,0 +1,187 @@
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from google.cloud import texttospeech
|
||||
|
||||
from aios import AIStorage,ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeNode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
"""
|
||||
You need to set the GOOGLE_APPLICATION_CREDENTIALS environment variable when using it.
|
||||
see:https://cloud.google.com/text-to-speech/docs/before-you-begin
|
||||
"""
|
||||
|
||||
|
||||
class GoogleTextToSpeechNode(ComputeNode):
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.node_id = "google_text_to_speech_node"
|
||||
self.task_queue = Queue()
|
||||
self.client: Optional[texttospeech.TextToSpeechClient] = None
|
||||
|
||||
self.language_list = {
|
||||
"cnm-CN": {
|
||||
"female": ["cmn-CN-Standard-A",
|
||||
"cmn-CN-Standard-D",
|
||||
"cmn-CN-Wavenet-A",
|
||||
"cmn-CN-Wavenet-D",
|
||||
"cmn-TW-Standard-A",
|
||||
"cmn-TW-Wavenet-A"],
|
||||
"man": ["cmn-CN-Standard-B",
|
||||
"cmn-CN-Standard-C",
|
||||
"cmn-CN-Wavenet-B",
|
||||
"cmn-CN-Wavenet-C",
|
||||
"cmn-TW-Standard-B",
|
||||
"cmn-TW-Standard-C",
|
||||
"cmn-TW-Wavenet-B",
|
||||
"cmn-TW-Wavenet-C"]
|
||||
},
|
||||
"en-US": {
|
||||
"female": ["en-US-Neural2-C",
|
||||
"en-US-Neural2-E",
|
||||
"en-US-Neural2-F",
|
||||
"en-US-Neural2-G",
|
||||
"en-US-Neural2-H",
|
||||
"en-US-News-K",
|
||||
"en-US-News-L",
|
||||
"en-US-Standard-C",
|
||||
"en-US-Standard-E",
|
||||
"en-US-Standard-F",
|
||||
"en-US-Standard-G",
|
||||
"en-US-Standard-H",
|
||||
"en-US-Studio-O",
|
||||
"en-US-Wavenet-C",
|
||||
"en-US-Wavenet-E",
|
||||
"en-US-Wavenet-F",
|
||||
"en-US-Wavenet-G",
|
||||
"en-US-Wavenet-H"],
|
||||
"man": ["en-US-Polyglot-1",
|
||||
"en-US-Standard-A",
|
||||
"en-US-Standard-B",
|
||||
"en-US-Standard-D",
|
||||
"en-US-Standard-I",
|
||||
"en-US-Standard-J",
|
||||
"en-US-Studio-M",
|
||||
"en-US-Wavenet-A",
|
||||
"en-US-Wavenet-B",
|
||||
"en-US-Wavenet-D",
|
||||
"en-US-Wavenet-I",
|
||||
"en-US-Wavenet-J"]
|
||||
}
|
||||
}
|
||||
self.start()
|
||||
|
||||
def init(self):
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
google_application_credentials = user_config.get_value("google_application_credentials")
|
||||
if google_application_credentials is None:
|
||||
raise Exception("google_application_credentials is None!")
|
||||
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = google_application_credentials
|
||||
self.client = texttospeech.TextToSpeechClient()
|
||||
|
||||
def start(self):
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
try:
|
||||
result = self._run_task(task)
|
||||
if result is not None:
|
||||
task.state = ComputeTaskState.DONE
|
||||
task.result = result
|
||||
except Exception as e:
|
||||
logger.error(f"google_text_to_speech_node run task error: {e}")
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.result = ComputeTaskResult()
|
||||
task.result.set_from_task(task)
|
||||
task.result.worker_id = self.node_id
|
||||
task.result.result_str = str(e)
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
def _run_task(self, task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
language_code = task.params["language_code"]
|
||||
if language_code is None:
|
||||
language_code = "en"
|
||||
|
||||
text = task.params["text"]
|
||||
voice_name = task.params["voice_name"]
|
||||
if voice_name is None:
|
||||
voice_name = "default"
|
||||
|
||||
gender = task.params["gender"]
|
||||
if gender is None:
|
||||
gender = "female"
|
||||
|
||||
age = task.params["age"]
|
||||
|
||||
if language_code == "zh":
|
||||
language_code = "cnm-CN"
|
||||
elif language_code == "en":
|
||||
language_code = "en-US"
|
||||
else:
|
||||
raise Exception(f"language_code {language_code} not support")
|
||||
|
||||
lang_list = self.language_list[language_code][gender]
|
||||
voice = lang_list[hash(voice_name) % len(lang_list)]
|
||||
|
||||
synthesis_input = texttospeech.SynthesisInput(text=text)
|
||||
voice = texttospeech.VoiceSelectionParams(language_code=language_code,
|
||||
ssml_gender=texttospeech.SsmlVoiceGender.NEUTRAL,
|
||||
name=voice)
|
||||
|
||||
audio_config = texttospeech.AudioConfig(audio_encoding=texttospeech.AudioEncoding.MP3)
|
||||
|
||||
response = self.client.synthesize_speech(input=synthesis_input, voice=voice, audio_config=audio_config)
|
||||
|
||||
result = ComputeTaskResult()
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
result.result = response.audio_content
|
||||
return result
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"google_text_to_speech_node push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
def display(self) -> str:
|
||||
return f"GoogleTextToSpeechNode: {self.node_id}"
|
||||
|
||||
def get_capacity(self):
|
||||
return 0
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
if task.task_type == ComputeTaskType.TEXT_2_VOICE:
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return False
|
||||
|
||||
def declare_user_config(self,is_optional:bool = False):
|
||||
if os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is None:
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
user_config.add_user_config("google_application_credentials",
|
||||
"google application credentials, please visit:https://cloud.google.com/text-to-speech/docs/before-you-begin",
|
||||
True,
|
||||
None)
|
||||
@@ -2,7 +2,7 @@ import os
|
||||
import runpy
|
||||
import toml
|
||||
import asyncio
|
||||
from knowledge import KnowledgePipelineEnvironment, KnowledgePipeline
|
||||
from aios import KnowledgePipelineEnvironment, KnowledgePipeline
|
||||
|
||||
|
||||
class KnowledgePipelineManager:
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .local_llama_compute_node import LocalLlama_ComputeNode
|
||||
@@ -0,0 +1,190 @@
|
||||
|
||||
import json
|
||||
import logging
|
||||
import requests
|
||||
from typing import Optional, List
|
||||
from pydantic import BaseModel
|
||||
|
||||
from aios import ComputeTask,Queue_ComputeNode, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType,AIStorage,UserConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
"""
|
||||
This is a custom implementation, it should be redesigned.
|
||||
"""
|
||||
|
||||
class LocalLlama_ComputeNode(Queue_ComputeNode):
|
||||
def __init__(self, url: str, model_name: str):
|
||||
super().__init__()
|
||||
self.url = url
|
||||
self.model_name = model_name
|
||||
|
||||
async def execute_task(self, task: ComputeTask)->ComputeTaskResult:
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
match task.task_type:
|
||||
case ComputeTaskType.TEXT_EMBEDDING:
|
||||
model_name = task.params["model_name"]
|
||||
input = task.params["input"]
|
||||
logger.info(f"call local-llama ({self.url}, {self.model_name}) {model_name} input: {input}")
|
||||
|
||||
self.embedding(input, result)
|
||||
|
||||
if result.result_code == ComputeTaskResultCode.OK:
|
||||
task.state = ComputeTaskState.DONE
|
||||
else:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = result.error_str
|
||||
|
||||
return result
|
||||
case ComputeTaskType.LLM_COMPLETION:
|
||||
mode_name = task.params["model_name"]
|
||||
prompts = task.params["prompts"]
|
||||
|
||||
logger.info(f"local-llama({self.url}, {self.model_name}) prompts: {prompts}")
|
||||
|
||||
self.completion(task, result)
|
||||
|
||||
if result.result_code == ComputeTaskResultCode.OK:
|
||||
task.state = ComputeTaskState.DONE
|
||||
else:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = result.error_str
|
||||
|
||||
case _:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
|
||||
result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
|
||||
return result
|
||||
|
||||
return result
|
||||
|
||||
async def initial(self) -> bool:
|
||||
return True
|
||||
|
||||
def display(self) -> str:
|
||||
return f"local-llama: {self.node_id}"
|
||||
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return True
|
||||
|
||||
def embedding(self, input: str, result: ComputeTaskResult):
|
||||
body = {
|
||||
"input": input
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(self.url + "/v1/embeddings", json = body, verify=False, headers={"Content-Type": "application/json"})
|
||||
response.close()
|
||||
|
||||
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
|
||||
|
||||
if response.status_code == 200:
|
||||
resp = response.json()
|
||||
result.result = resp["data"][0]["embedding"]
|
||||
elif response.status_code == 422:
|
||||
resp = response.json()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
|
||||
else:
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.error_str = "http request failed: " + str(response.status_code)
|
||||
except Exception as e:
|
||||
logger.error(f"call local-llama({self.url}, {self.model_name}) run TEXT_EMBEDDING task error: {e}")
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.error_str = str(e)
|
||||
return result
|
||||
|
||||
def completion(self, task: ComputeTask, result: ComputeTaskResult):
|
||||
mode_name = task.params["model_name"]
|
||||
prompts = task.params["prompts"]
|
||||
max_token_size = task.params.get("max_token_size")
|
||||
llm_inner_functions = task.params.get("inner_functions")
|
||||
if max_token_size is None:
|
||||
max_token_size = max_token_size
|
||||
|
||||
body = {
|
||||
"messages": [],
|
||||
"functions": llm_inner_functions,
|
||||
"tools": [],
|
||||
"tool_choices": [],
|
||||
"max_tokens": 4000
|
||||
}
|
||||
if llm_inner_functions is not None:
|
||||
for fun in llm_inner_functions:
|
||||
body["tools"].append({
|
||||
"type": "function",
|
||||
"function": fun,
|
||||
})
|
||||
body["tool_choices"].append({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": fun["name"]
|
||||
}
|
||||
})
|
||||
|
||||
for prompt in prompts:
|
||||
body["messages"].append({
|
||||
"role": prompt["role"],
|
||||
"content": prompt["content"]
|
||||
})
|
||||
|
||||
try:
|
||||
logger.info(f"will post http request to {self.url}/v1/chat/completions, body: {body}")
|
||||
|
||||
response = requests.post(self.url + "/v1/chat/completions", json = body, verify=False, headers={"Content-Type": "application/json"})
|
||||
response.close()
|
||||
|
||||
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
|
||||
|
||||
if response.status_code == 200:
|
||||
resp = response.json()
|
||||
|
||||
status_code = resp["choices"][0]["finish_reason"]
|
||||
token_usage = resp["usage"]
|
||||
|
||||
match status_code:
|
||||
case "tool_calls":
|
||||
task.state = ComputeTaskState.DONE
|
||||
# rebuild the function name
|
||||
fun_name = resp["choices"][0]["message"]["function_call"]["name"]
|
||||
if len(llm_inner_functions) == 1 and (fun_name is None or fun_name == ""):
|
||||
resp["choices"][0]["message"]["function_call"]["name"] = llm_inner_functions[0]["name"]
|
||||
case "stop":
|
||||
task.state = ComputeTaskState.DONE
|
||||
case _:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = f"The status code was {status_code}."
|
||||
result.error_str = f"The status code was {status_code}."
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
return None
|
||||
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.result_str = resp["choices"][0]["message"]["content"]
|
||||
result.result["message"] = resp["choices"][0]["message"]
|
||||
|
||||
if token_usage:
|
||||
result.result_refers["token_usage"] = token_usage
|
||||
|
||||
logger.info(f"local-llama({self.url}, {self.model_name}) success response: {result.result_str}")
|
||||
elif response.status_code == 422:
|
||||
resp = response.json()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
|
||||
else:
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.error_str = "http request failed: " + str(response.status_code)
|
||||
except Exception as e:
|
||||
logger.error(f"call local-llama({self.url}, {self.model_name}) run LLM_COMPLETION task error: {e}")
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.error_str = str(e)
|
||||
return result
|
||||
@@ -1,8 +1,7 @@
|
||||
# define a knowledge base class
|
||||
import json
|
||||
import string
|
||||
from aios_kernel import AIStorage, Environment, SimpleAIFunction, CustomAIAgent, AgentPrompt, AgentMsg
|
||||
from knowledge import *
|
||||
from aios import *
|
||||
from .mail import MailStorage, Mail
|
||||
|
||||
class IssueState(Enum):
|
||||
|
||||
@@ -2,8 +2,7 @@ import os
|
||||
import logging
|
||||
import json
|
||||
import string
|
||||
from knowledge import *
|
||||
from aios_kernel.storage import AIStorage
|
||||
from aios import *
|
||||
from .mail import Mail, MailStorage
|
||||
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ import datetime
|
||||
from bs4 import BeautifulSoup
|
||||
import sqlite3
|
||||
import html2text
|
||||
from knowledge import *
|
||||
from aios import *
|
||||
|
||||
class Mail:
|
||||
def __init__(self, **kwargs) -> None:
|
||||
|
||||
@@ -3,8 +3,7 @@ import logging
|
||||
import json
|
||||
import imaplib
|
||||
import mailparser
|
||||
from knowledge import *
|
||||
from aios_kernel.storage import AIStorage
|
||||
from aios import *
|
||||
|
||||
|
||||
class EmailSpider:
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
from .cid import ContentId
|
||||
from .ndn_client import NDN_Client
|
||||
@@ -1,12 +0,0 @@
|
||||
|
||||
|
||||
class ContentId:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def as_str(self) -> str:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def create_from_str(cid_str:str):
|
||||
pass
|
||||
@@ -1,117 +0,0 @@
|
||||
import asyncio,aiofiles,aiohttp
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from .cid import ContentId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NDN_GET_TASK_STATE_INIT = 0
|
||||
NDN_GET_TAKS_CONNECTING = 1
|
||||
NDN_GET_TASK_STATE_DOWNLOADING = 2
|
||||
NDN_GET_TASK_STATE_VERIFYING = 3
|
||||
NDN_GET_TASK_STATE_DONE = 4
|
||||
NDN_GET_TASK_STATE_ERROR = 5
|
||||
|
||||
class NDN_GetTask:
|
||||
def __init__(self) -> None:
|
||||
self.cid:str = None
|
||||
self.target_path:str = None
|
||||
self.urls:[str] = None
|
||||
self.options:Optional[dict] = None
|
||||
|
||||
self.working_task = None
|
||||
self.state = NDN_GET_TASK_STATE_INIT
|
||||
self.total_size = 0
|
||||
self.recv_bytes = 0
|
||||
self.write_bytes = 0
|
||||
self.error_str = None
|
||||
self.chunk_queue = None
|
||||
|
||||
self.retry_count = 0
|
||||
self.used_urls = []
|
||||
self.hash_update = None
|
||||
|
||||
|
||||
def select_url(self,index:int)->str:
|
||||
return self.urls[0]
|
||||
|
||||
def get_chunk_for_download(self)->bytes:
|
||||
pass
|
||||
|
||||
class NDN_Client:
|
||||
def __init__(self):
|
||||
self.cache_dir = ""
|
||||
self.default_ndn_http_gateway = ""
|
||||
self.all_task = {}
|
||||
self.memory_chunk_size = 1024*1024*2
|
||||
self.chunk_queue_size = 16
|
||||
|
||||
def load_config(self,config:dict):
|
||||
if config.get("cache_dir"):
|
||||
self.cache_dir = config.get("cache_dir")
|
||||
if config.get("dndn_gateway"):
|
||||
self.default_ndn_http_gateway = config.get("ndn_gateway")
|
||||
|
||||
def get_file(self,cid:ContentId,target_path:str,urls:{}=None,options:{}=None)->NDN_GetTask:
|
||||
get_task = self.all_task.get(cid.as_str())
|
||||
if get_task:
|
||||
return get_task
|
||||
else:
|
||||
get_task = NDN_GetTask()
|
||||
self.all_task[cid.as_str()] = get_task
|
||||
|
||||
get_task.cid = cid
|
||||
get_task.target_path = target_path
|
||||
get_task.urls = urls
|
||||
get_task.options = options
|
||||
if get_task.urls is None:
|
||||
get_task.urls = [f"{self.default_ndn_http_gateway}/{cid.as_str()}"]
|
||||
logger.info(f"get_file {cid.as_str()} urls is None, use {get_task.urls[0]} as default")
|
||||
|
||||
|
||||
async def get_file_async():
|
||||
target_file = aiofiles.open(target, 'wb')
|
||||
# if file exist, check hash first
|
||||
|
||||
http_session = aiohttp.ClientSession()
|
||||
resp = http_session.get(get_task.select_url(0))
|
||||
if resp.status != 200:
|
||||
get_task.error_str = f"get_file {cid.as_str()} failed,http status:{resp.status}"
|
||||
return
|
||||
get_task.total_size = resp.content_length
|
||||
|
||||
async def write_file_async():
|
||||
while True:
|
||||
chunk = await get_task.chunk_queue.pop()
|
||||
chunk_size = len(chunk)
|
||||
if not chunk or chunk_size == 0:
|
||||
break
|
||||
get_task.hash_update.update(chunk)
|
||||
await target_file.write(chunk)
|
||||
get_task.write_bytes += chunk_size
|
||||
|
||||
#verify
|
||||
get_task.state = NDN_GET_TASK_STATE_VERIFYING
|
||||
await target_file.close()
|
||||
return
|
||||
|
||||
write_task = asyncio.create_task(write_file_async())
|
||||
while True:
|
||||
await get_task.chunk_queue.pop()
|
||||
chunk = resp.content.read(self.memory_chunk_size)
|
||||
chunk_size = len(chunk)
|
||||
if not chunk or chunk_size == 0:
|
||||
break
|
||||
|
||||
get_task.recv_bytes += len(chunk)
|
||||
get_task.chunk_queue.push(chunk)
|
||||
|
||||
|
||||
get_task.state = NDN_GET_TASK_STATE_DONE
|
||||
await write_task
|
||||
|
||||
get_task.working_task = asyncio.create_task(get_file_async())
|
||||
return get_task
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from .open_ai_node import *
|
||||
from .openai_tts_node import *
|
||||
from .whisper_node import *
|
||||
@@ -0,0 +1,301 @@
|
||||
import openai
|
||||
from openai import AsyncOpenAI
|
||||
import os
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
import json
|
||||
import aiohttp
|
||||
import base64
|
||||
import requests
|
||||
|
||||
from aios import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeTaskResultCode,ComputeNode,AIStorage,UserConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OpenAI_ComputeNode(ComputeNode):
|
||||
_instance = None
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = OpenAI_ComputeNode()
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
def declare_user_config(cls):
|
||||
if os.getenv("OPENAI_API_KEY_") is None:
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
user_config.add_user_config("openai_api_key","openai api key",False,None)
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.is_start = False
|
||||
# openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud
|
||||
self.openai_api_key = None
|
||||
self.node_id = "openai_node"
|
||||
self.task_queue = Queue()
|
||||
|
||||
|
||||
async def initial(self):
|
||||
if os.getenv("OPENAI_API_KEY") is not None:
|
||||
self.openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
else:
|
||||
self.openai_api_key = AIStorage.get_instance().get_user_config().get_value("openai_api_key")
|
||||
|
||||
if self.openai_api_key is None:
|
||||
logger.error("openai_api_key is None!")
|
||||
return False
|
||||
|
||||
openai.api_key = self.openai_api_key
|
||||
self.start()
|
||||
return True
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"openai_node push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
def message_to_dict(self, message)->dict:
|
||||
result = message.dict()
|
||||
# result_msg = {}
|
||||
# #message.json()
|
||||
# if message.content:
|
||||
# result_msg["content"] = message.content
|
||||
# result_msg["role"] = message.role
|
||||
# if message.function_call:
|
||||
# function_call = {}
|
||||
# function_call["arguments"] = message.function_call.arguments
|
||||
# function_call["name"] = message.function_call.name
|
||||
# result_msg["function_call"] = function_call
|
||||
|
||||
# if message.tool_calls:
|
||||
# tool_calls = []
|
||||
# for tool_call in message.tool_calls:
|
||||
# tool_call_dict = {}
|
||||
# tool_call_dict["id"] = tool_call.id
|
||||
# tool_call_dict["type"] = tool_call.type
|
||||
# func_call_dict = {}
|
||||
# func_call_dict["name"] = tool_call.function.name
|
||||
# func_call_dict["arguments"] = tool_call.function.arguments
|
||||
# tool_call_dict["function"] = func_call_dict
|
||||
|
||||
# tool_calls.append(tool_call_dict)
|
||||
# result_msg["tool_calls"] = message.tool_calls
|
||||
|
||||
# result["message"] = result_msg
|
||||
return result
|
||||
|
||||
def _image_2_text(self, task: ComputeTask):
|
||||
logger.info('openai image_2_text')
|
||||
# 本地图片处理
|
||||
def encode_image(image_path):
|
||||
with open(image_path, "rb") as image_file:
|
||||
return base64.b64encode(image_file.read()).decode('utf-8')
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.openai_api_key }"
|
||||
}
|
||||
model_name = task.params["model_name"]
|
||||
base64_image = encode_image(task.params["image_path"])
|
||||
payload = {
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": task.params["prompt"]
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"max_tokens": 300
|
||||
}
|
||||
logger.info('openai send image_2_text request ')
|
||||
# openai 的库的Vision只支持传图片的url地址。本地图片得用request
|
||||
response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
|
||||
if response.status_code == 200:
|
||||
logger.info('openai image_2_text success')
|
||||
return response.json()
|
||||
else:
|
||||
logger.error('openai image_2_text error')
|
||||
logger.error(response.json())
|
||||
return None
|
||||
|
||||
async def _run_task(self, task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
|
||||
match task.task_type:
|
||||
case ComputeTaskType.TEXT_EMBEDDING:
|
||||
model_name = task.params["model_name"]
|
||||
input = task.params["input"]
|
||||
logger.info(f"call openai {model_name} input: {input}")
|
||||
try:
|
||||
resp = openai.Embedding.create(model=model_name,
|
||||
input=input)
|
||||
except Exception as e:
|
||||
logger.error(f"openai run TEXT_EMBEDDING task error: {e}")
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = str(e)
|
||||
result.error_str = str(e)
|
||||
return result
|
||||
|
||||
# resp = {
|
||||
# "object": "list",
|
||||
# "data": [
|
||||
# {
|
||||
# "object": "embedding",
|
||||
# "index": 0,
|
||||
# "embedding": [
|
||||
# -0.00930514745414257,
|
||||
# 0.00765434792265296,
|
||||
# -0.007167573552578688,
|
||||
# -0.012373941019177437,
|
||||
# -0.04884673282504082
|
||||
# ]}]
|
||||
# }
|
||||
|
||||
logger.info(f"openai response: {resp}")
|
||||
task.state = ComputeTaskState.DONE
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.worker_id = self.node_id
|
||||
result.result_str = resp["data"][0]["embedding"]
|
||||
|
||||
return result
|
||||
case ComputeTaskType.IMAGE_2_TEXT:
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.worker_id = self.node_id
|
||||
# result.result_str = resp["data"][0]["image_2_text"]
|
||||
result.result["message"] = self._image_2_text(task)
|
||||
return result
|
||||
case ComputeTaskType.LLM_COMPLETION:
|
||||
mode_name = task.params["model_name"]
|
||||
prompts = task.params["prompts"]
|
||||
resp_mode = task.params["resp_mode"]
|
||||
if resp_mode == "json":
|
||||
response_format = { "type": "json_object" }
|
||||
else:
|
||||
response_format = None
|
||||
max_token_size = task.params.get("max_token_size")
|
||||
llm_inner_functions = task.params.get("inner_functions")
|
||||
if max_token_size is None:
|
||||
max_token_size = 4000
|
||||
|
||||
result_token = max_token_size
|
||||
client = AsyncOpenAI(api_key=self.openai_api_key)
|
||||
try:
|
||||
if llm_inner_functions is None:
|
||||
logger.info(f"call openai {mode_name} prompts: {prompts}")
|
||||
resp = await client.chat.completions.create(model=mode_name,
|
||||
messages=prompts,
|
||||
response_format = response_format,
|
||||
#max_tokens=result_token,
|
||||
)
|
||||
else:
|
||||
logger.info(f"call openai {mode_name} prompts: \n\t {prompts} \nfunctions: \n\t{json.dumps(llm_inner_functions)}")
|
||||
resp = await client.chat.completions.create(model=mode_name,
|
||||
messages=prompts,
|
||||
response_format = response_format,
|
||||
functions=llm_inner_functions,
|
||||
# max_tokens=result_token,
|
||||
) # TODO: add temperature to task params?
|
||||
except Exception as e:
|
||||
logger.error(f"openai run LLM_COMPLETION task error: {e}")
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = str(e)
|
||||
result.error_str = str(e)
|
||||
return result
|
||||
|
||||
logger.info(f"openai response: {resp}")
|
||||
status_code = resp.choices[0].finish_reason
|
||||
token_usage = resp.usage
|
||||
|
||||
match status_code:
|
||||
case "function_call":
|
||||
task.state = ComputeTaskState.DONE
|
||||
case "stop":
|
||||
task.state = ComputeTaskState.DONE
|
||||
case _:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = f"The status code was {status_code}."
|
||||
result.error_str = f"The status code was {status_code}."
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
return result
|
||||
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.worker_id = self.node_id
|
||||
result.result_str = resp.choices[0].message.content
|
||||
|
||||
result.result["message"] = self.message_to_dict(resp.choices[0].message)
|
||||
|
||||
if token_usage:
|
||||
result.result_refers["token_usage"] = token_usage
|
||||
logger.info(f"openai success response: {result.result_str}")
|
||||
return result
|
||||
case _:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
|
||||
result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
|
||||
return None
|
||||
|
||||
def start(self):
|
||||
if self.is_start is True:
|
||||
return
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
logger.info(f"openai_node get task: {task.display()}")
|
||||
result = await self._run_task(task)
|
||||
if result is not None:
|
||||
task.result = result
|
||||
task.state = ComputeTaskState.DONE
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
def display(self) -> str:
|
||||
return f"OpenAI_ComputeNode: {self.node_id}"
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
if task.task_type == ComputeTaskType.LLM_COMPLETION:
|
||||
if not task.params["model_name"]:
|
||||
return True
|
||||
model_name : str = task.params["model_name"]
|
||||
if model_name.startswith("gpt-"):
|
||||
return True
|
||||
|
||||
if task.task_type == ComputeTaskType.IMAGE_2_TEXT:
|
||||
model_name : str = task.params["model_name"]
|
||||
if model_name.startswith("gpt-4"):
|
||||
return True
|
||||
#if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
|
||||
# if task.params["model_name"] == "text-embedding-ada-002":
|
||||
# return True
|
||||
return False
|
||||
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return False
|
||||
@@ -0,0 +1,118 @@
|
||||
import asyncio
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
from asyncio import Queue
|
||||
|
||||
from aios import ComputeNode, ComputeTask, ComputeTaskState, ComputeTaskResult, ComputeTaskType, AIStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OpenAITTSComputeNode(ComputeNode):
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.is_start = False
|
||||
self.node_id = "openai_tts_node"
|
||||
self.task_queue = Queue()
|
||||
self.voice_list = {
|
||||
"female": ["nova", "shimmer"],
|
||||
"man": ["alloy", "echo", "fable", "onyx"]
|
||||
}
|
||||
if os.getenv("OPENAI_API_KEY") is not None:
|
||||
self.openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
else:
|
||||
self.openai_api_key = AIStorage.get_instance().get_user_config().get_value("openai_api_key")
|
||||
|
||||
self.start()
|
||||
|
||||
def start(self):
|
||||
if self.is_start is True:
|
||||
logger.warn("OpenAITTSComputeNode is already start")
|
||||
return
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
try:
|
||||
result = await self._run_task(task)
|
||||
if result is not None:
|
||||
task.state = ComputeTaskState.DONE
|
||||
task.result = result
|
||||
except Exception as e:
|
||||
logger.error(f"openai_tts_node run task error: {e}")
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.result = ComputeTaskResult()
|
||||
task.result.set_from_task(task)
|
||||
task.result.worker_id = self.node_id
|
||||
task.result.result_str = str(e)
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
async def _run_task(self,task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
text = task.params["text"]
|
||||
voice_name = task.params["voice_name"]
|
||||
if voice_name is None:
|
||||
voice_name = "default"
|
||||
gender = task.params["gender"]
|
||||
if gender is None:
|
||||
gender = "female"
|
||||
|
||||
voice_list = self.voice_list[gender]
|
||||
voice = voice_list[hash(voice_name)%len(voice_list)]
|
||||
|
||||
model_name = task.params['model_name']
|
||||
if model_name is None:
|
||||
model_name = 'tts-1'
|
||||
|
||||
client = AsyncOpenAI(api_key=self.openai_api_key)
|
||||
|
||||
response = await client.audio.speech.create(model=model_name, voice=voice, input=text)
|
||||
|
||||
cache = io.BytesIO()
|
||||
async for data in await response.aiter_bytes():
|
||||
cache.write(data)
|
||||
|
||||
cache.seek(0)
|
||||
|
||||
result = ComputeTaskResult()
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
result.result = cache.read()
|
||||
return result
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"openai_tts_node push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
def display(self) -> str:
|
||||
return f"OpenAITTSComputeNode: {self.node_id}"
|
||||
|
||||
def get_capacity(self):
|
||||
return 0
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
if task.task_type == ComputeTaskType.TEXT_2_VOICE:
|
||||
if task.params['model_name'] is None or task.params['model_name'] == 'tts-1' or task.params['model_name'] == 'tts-1-hd':
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return False
|
||||
@@ -0,0 +1,226 @@
|
||||
import io
|
||||
import json
|
||||
from asyncio import Queue
|
||||
import asyncio
|
||||
import openai
|
||||
import os
|
||||
import logging
|
||||
import srt
|
||||
import webvtt
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from openai.cli._progress import BufferReader
|
||||
from pydub import AudioSegment
|
||||
from datetime import timedelta
|
||||
|
||||
from aios import AIStorage,ComputeNode,ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SECONDS_IN_HOUR = 3600
|
||||
SECONDS_IN_MINUTE = 60
|
||||
HOURS_IN_DAY = 24
|
||||
MICROSECONDS_IN_MILLISECOND = 1000
|
||||
|
||||
def timedelta_to_vtt_timestamp(timedelta_timestamp):
|
||||
hrs, secs_remainder = divmod(timedelta_timestamp.seconds, SECONDS_IN_HOUR)
|
||||
hrs += timedelta_timestamp.days * HOURS_IN_DAY
|
||||
mins, secs = divmod(secs_remainder, SECONDS_IN_MINUTE)
|
||||
msecs = timedelta_timestamp.microseconds // MICROSECONDS_IN_MILLISECOND
|
||||
return "%02d:%02d:%02d.%03d" % (hrs, mins, secs, msecs)
|
||||
|
||||
|
||||
class WhisperComputeNode(ComputeNode):
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.is_start = False
|
||||
self.node_id = "whisper_node"
|
||||
self.enable = True
|
||||
self.task_queue = Queue()
|
||||
|
||||
if os.getenv("OPENAI_API_KEY") is not None:
|
||||
self.openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
else:
|
||||
self.openai_api_key = AIStorage.get_instance().get_user_config().get_value("openai_api_key")
|
||||
|
||||
self.start()
|
||||
|
||||
def start(self):
|
||||
if self.is_start is True:
|
||||
logger.warn("WhisperComputeNode is already start")
|
||||
return
|
||||
self.is_start = True
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
try:
|
||||
result = await self._run_task(task)
|
||||
if result is not None:
|
||||
task.state = ComputeTaskState.DONE
|
||||
task.result = result
|
||||
except Exception as e:
|
||||
logger.error(f"whisper_node run task error: {e}")
|
||||
logger.exception(e)
|
||||
task.state = ComputeTaskState.ERROR
|
||||
task.result = ComputeTaskResult()
|
||||
task.result.set_from_task(task)
|
||||
task.result.worker_id = self.node_id
|
||||
task.result.result_str = str(e)
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
async def _run_task(self, task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
prompt = task.params["prompt"]
|
||||
response_format = None
|
||||
if "response_format" in task.params:
|
||||
response_format = task.params["response_format"]
|
||||
temperature = None
|
||||
if "temperature" in task.params:
|
||||
temperature = task.params["temperature"]
|
||||
language = None
|
||||
if "language" in task.params:
|
||||
language = task.params["language"]
|
||||
file = task.params["file"]
|
||||
|
||||
client = AsyncOpenAI(api_key=self.openai_api_key)
|
||||
|
||||
if os.path.getsize(file) > 25 * 1024 * 1024:
|
||||
audio = AudioSegment.from_file(file)
|
||||
text = ""
|
||||
results = []
|
||||
latest_resp = None
|
||||
step = 10 * 60 * 1000
|
||||
for i in range(0, len(audio), step):
|
||||
if i + step < len(audio):
|
||||
chunk = audio[i:i + step]
|
||||
else:
|
||||
chunk = audio[i:]
|
||||
seg_file = io.BytesIO()
|
||||
chunk.export(seg_file, format="mp3")
|
||||
seg_file.seek(0)
|
||||
|
||||
resp = await client.audio.transcriptions.create(model="whisper-1",
|
||||
file = ("test.mp3", seg_file),
|
||||
language=language,
|
||||
temperature=temperature,
|
||||
prompt=prompt,
|
||||
response_format=response_format)
|
||||
if response_format == "json":
|
||||
if text == "":
|
||||
text = resp.text
|
||||
else:
|
||||
text += "," + resp.text
|
||||
elif response_format == "text":
|
||||
if text == "":
|
||||
text = resp
|
||||
else:
|
||||
text += "," + resp
|
||||
elif response_format == "verbose_json":
|
||||
if text == "":
|
||||
text = resp.text
|
||||
else:
|
||||
text += "," + resp.text
|
||||
results.extend(resp.segments)
|
||||
elif response_format == "srt":
|
||||
srt_list = list(srt.parse(resp))
|
||||
for item in srt_list:
|
||||
item.start += timedelta(milliseconds=i)
|
||||
item.end += timedelta(milliseconds=i)
|
||||
results.append(item)
|
||||
elif response_format == "vtt":
|
||||
vtt = webvtt.read_buffer(io.StringIO(resp))
|
||||
for caption in vtt.captions:
|
||||
start = timedelta_to_vtt_timestamp(
|
||||
srt.srt_timestamp_to_timedelta(caption.start) + timedelta(milliseconds=i))
|
||||
end = timedelta_to_vtt_timestamp(
|
||||
srt.srt_timestamp_to_timedelta(caption.end) + timedelta(milliseconds=i))
|
||||
results.append(webvtt.Caption(start, end, caption.text))
|
||||
else:
|
||||
raise Exception(f"not support response_format: {response_format}")
|
||||
|
||||
latest_resp = resp
|
||||
|
||||
result = ComputeTaskResult()
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
if response_format == "text":
|
||||
result.result_str = text
|
||||
result.result = text
|
||||
elif response_format == "json":
|
||||
result.result_str = json.dumps({"text": text})
|
||||
resp.text = text
|
||||
result.result = resp
|
||||
elif response_format == "verbose_json":
|
||||
result.result_str = json.dumps({"text": text, "segments": results})
|
||||
latest_resp.text = text
|
||||
latest_resp.segments = results
|
||||
result.result = latest_resp
|
||||
elif response_format == "srt":
|
||||
result.result_str = srt.compose(results)
|
||||
result.result = result.result_str
|
||||
elif response_format == "vtt":
|
||||
vtt = webvtt.WebVTT()
|
||||
vtt.captions.extend(results)
|
||||
f = io.StringIO()
|
||||
vtt.write(f)
|
||||
f.seek(0)
|
||||
result.result_str = f.read()
|
||||
result.result = result.result_str
|
||||
return result
|
||||
else:
|
||||
with open(file, "rb") as file_reader:
|
||||
buffer_reader = BufferReader(file_reader.read(), desc="Upload progress")
|
||||
|
||||
resp = await client.audio.transcriptions.create(model="whisper-1",
|
||||
file = (file, buffer_reader),
|
||||
language=language,
|
||||
temperature=temperature,
|
||||
prompt=prompt,
|
||||
response_format=response_format)
|
||||
result = ComputeTaskResult()
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
if response_format == "json":
|
||||
result.result_str = json.dumps({"text": resp.text})
|
||||
elif response_format == "verbose_json":
|
||||
result.result_str = json.dumps({"text": resp.text, "segments": resp.segments})
|
||||
elif response_format == "srt" or response_format == "vtt" or response_format == "text":
|
||||
result.result_str = resp
|
||||
else:
|
||||
raise Exception(f"not support response_format: {response_format}")
|
||||
result.result = resp
|
||||
return result
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"whisper_node push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
def display(self) -> str:
|
||||
return f"WhisperComputeNode: {self.node_id}"
|
||||
|
||||
def get_capacity(self):
|
||||
return 0
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
if task.task_type == ComputeTaskType.VOICE_2_TEXT:
|
||||
if task.params['model_name'] is None or task.params['model_name'] == 'openai-whisper':
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return False
|
||||
@@ -1 +0,0 @@
|
||||
TODO
|
||||
@@ -1,3 +0,0 @@
|
||||
from .env import PackageEnvManager,PackageEnv
|
||||
from .pkg import PackageInfo,PackageMediaInfo
|
||||
from .installer import PackageInstallTask
|
||||
@@ -1,158 +0,0 @@
|
||||
|
||||
import logging
|
||||
import toml
|
||||
import os
|
||||
|
||||
from .pkg import PackageInfo,PackageMediaInfo
|
||||
from .media_reader import MediaReader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PackageEnv:
|
||||
def __init__(self,cfg_path:str) -> None:
|
||||
self.pkg_dir : str = "./pkgs/"
|
||||
self.pkg_obj_dir : str = "./.pkgs/"
|
||||
|
||||
self.locked_index : str = "./pkg.lock"
|
||||
self.is_strict : bool = True
|
||||
self.parent_envs : list[PackageEnv] = []
|
||||
self.index_dbs = None
|
||||
|
||||
self.env_dir = None
|
||||
self.cfg_path = cfg_path
|
||||
self._load_pkg_cfg(cfg_path)
|
||||
pass
|
||||
|
||||
def load_from_config(self,config:dict) -> bool:
|
||||
if config.get("main") is not None:
|
||||
self.pkg_dir = os.path.abspath(self.env_dir + "/" + config["main"])
|
||||
|
||||
if config.get("cache") is not None:
|
||||
self.pkg_obj_dir = os.path.abspath(self.env_dir + "/ " + config["cache"])
|
||||
|
||||
def load(self,pkg_name:str,search_parent=True) -> PackageMediaInfo:
|
||||
pkg_path = None
|
||||
pkg_id,verion_str,cid = PackageInfo.parse_pkg_name(pkg_name)
|
||||
|
||||
if cid is None:
|
||||
if verion_str is None:
|
||||
pkg_path = f"{self.pkg_dir}/{pkg_id}"
|
||||
else:
|
||||
#TODO fix bug about channel here
|
||||
channel:str = self.get_pkg_channel_from_version(verion_str)
|
||||
the_version:str = self.get_exact_version_from_installed(verion_str)
|
||||
if the_version is None:
|
||||
logger.warn(f"load {pkg_name} failed: no match version from {verion_str}")
|
||||
return None
|
||||
if channel is None:
|
||||
pkg_path = f"{self.pkg_dir}/{pkg_id}#{the_version}"
|
||||
else:
|
||||
pkg_path = f"{self.pkg_dir}/{pkg_id}#{channel}#{the_version}"
|
||||
else:
|
||||
pkg_path = f"{self.pkg_obj_dir}/.{pkg_id}/{cid}"
|
||||
|
||||
media_info:PackageMediaInfo = self.try_load_pkg_media_info(pkg_path)
|
||||
if media_info is None:
|
||||
if search_parent is True and self.parent_envs is not None:
|
||||
for parent_env in self.parent_envs:
|
||||
media_info = parent_env.load(pkg_id,False)
|
||||
if media_info is not None:
|
||||
return media_info
|
||||
|
||||
if media_info is None:
|
||||
logger.warn(f"pkg_load {pkg_id}, cid:{cid} error,not found ,search_parent={search_parent}")
|
||||
|
||||
return media_info
|
||||
|
||||
def get_exact_version_from_installed(self,verion_str:str) -> str:
|
||||
pass
|
||||
|
||||
def get_pkg_channel_from_version(self,pkg_version:str) -> str:
|
||||
args = pkg_version.split("~")
|
||||
if len(args) == 1:
|
||||
return None
|
||||
else:
|
||||
return args[0]
|
||||
|
||||
|
||||
def get_pkg_media_info(self,pkg_name:str)->PackageMediaInfo:
|
||||
pass
|
||||
|
||||
def try_load_pkg_media_info(self,pkg_full_path:str) -> PackageMediaInfo:
|
||||
the_result : PackageMediaInfo = None
|
||||
logger.debug(f"try load pkng from:{pkg_full_path}")
|
||||
if os.path.isdir(pkg_full_path):
|
||||
the_result = PackageMediaInfo(pkg_full_path,"dir")
|
||||
|
||||
return the_result
|
||||
|
||||
def _create_media_loader(self,media_info:PackageMediaInfo) -> MediaReader:
|
||||
match media_info.media_type:
|
||||
case "dir":
|
||||
from .media_reader import FolderMediaReader
|
||||
return FolderMediaReader(media_info.full_path)
|
||||
|
||||
logger.error(f"create media loader for {media_info} failed!")
|
||||
return None
|
||||
|
||||
def get_installed_pkg_info(self,pkg_name:str) -> PackageInfo:
|
||||
pass
|
||||
|
||||
def lookup(self,pkg_id:str,version_str:str) -> PackageInfo:
|
||||
# to make sure pkg.cid is correct, we MUST verfiy eveything here
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def is_valied_media(pkg_full_path:str) -> bool:
|
||||
pass
|
||||
|
||||
def do_pkg_media_trans(self,pkg_info:PackageInfo,source_path:str,target_path:str) -> bool:
|
||||
pass
|
||||
|
||||
def _load_pkg_cfg(self,cfg_path:str):
|
||||
if cfg_path is None:
|
||||
return
|
||||
|
||||
cfg = None
|
||||
if len(cfg_path) < 1:
|
||||
return
|
||||
try:
|
||||
cfg = toml.load(cfg_path)
|
||||
self.env_dir = os.path.abspath(os.path.dirname(cfg_path))
|
||||
self.cfg_path = os.path.abspath(cfg_path)
|
||||
except Exception as e:
|
||||
logger.error(f"read pkg cfg from {cfg_path} failed! unexpected error occurred: {str(e)}")
|
||||
return
|
||||
|
||||
return self.load_from_config(cfg)
|
||||
|
||||
|
||||
|
||||
def _preprocess_prefixs(self,prefixs):
|
||||
pass
|
||||
|
||||
class PackageEnvManager:
|
||||
_instance = None
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = PackageEnvManager()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._pkg_envs = {}
|
||||
|
||||
def get_env(self,cfg_path:str) -> PackageEnv:
|
||||
if cfg_path in self._pkg_envs:
|
||||
return self._pkg_envs[cfg_path]
|
||||
else:
|
||||
pkg_env = PackageEnv(cfg_path)
|
||||
self._pkg_envs[cfg_path] = pkg_env
|
||||
return pkg_env
|
||||
|
||||
def get_user_env(self) -> PackageEnv:
|
||||
pass
|
||||
|
||||
def get_system_env(self) -> PackageEnv:
|
||||
pass
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
@@ -1,170 +0,0 @@
|
||||
# installer download pkg by cid, than install it to target dir
|
||||
import logging
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import aiofiles
|
||||
import os
|
||||
from typing import Tuple
|
||||
|
||||
from ndn_client import ContentId,NDN_Client
|
||||
from .pkg import PackageInfo,PackageMediaInfo
|
||||
from .env import PackageEnv
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
INSTALL_TASK_STATE_DONE = 0
|
||||
INSTALL_TASK_STATE_CHECK_DEPENDENCY = 1
|
||||
INSTALL_TASK_STATE_INSTALL_DEPENDENCY = 2
|
||||
INSTALL_TASK_STATE_DOWNLOADING = 3
|
||||
INSTALL_TASK_STATE_INSTALLING = 4
|
||||
INSTALL_TAKS_STATE_ERROR = 5
|
||||
|
||||
class PackageInstallTask:
|
||||
def __init__(self,owner:PackageEnv) -> None:
|
||||
self.owner = owner
|
||||
self.state = INSTALL_TASK_STATE_CHECK_DEPENDENCY
|
||||
|
||||
self.pkg_media_info = None
|
||||
self.working_task = None
|
||||
self.dependency_tasks = None
|
||||
self.error_str = None
|
||||
|
||||
class PackageInstaller:
|
||||
def __init__(self,owner_env:PackageEnv) -> None:
|
||||
self.all_tasks = {}
|
||||
self.owner_env = owner_env
|
||||
|
||||
def install(self,pkg_name:str,
|
||||
install_from_dependency = False, can_upgrade = True,skip_depends = False,options = None)->Tuple[PackageInstallTask,str]:
|
||||
|
||||
the_pkg_info : PackageInfo = None
|
||||
is_upgrade : bool = False
|
||||
need_backup : bool = False
|
||||
|
||||
pkg_id,version_str,cid = PackageInfo.parse_pkg_name(pkg_name)
|
||||
media_info : PackageMediaInfo = self.owner_env.get_media_info(pkg_name) # must use index-db?
|
||||
if media_info is not None:
|
||||
if cid is not None:
|
||||
if can_upgrade:
|
||||
is_upgrade = True
|
||||
else:
|
||||
error_str = f"{pkg_name},{cid} already installed!"
|
||||
logger.error(error_str)
|
||||
return None,error_str
|
||||
else:
|
||||
the_pkg_info = self.owner_env.lookup(pkg_id,version_str,None)
|
||||
if the_pkg_info is None:
|
||||
error_str = f"{pkg_name} old version exist in local but not found in index db!"
|
||||
logger.error(error_str)
|
||||
return None,error_str
|
||||
else:
|
||||
is_upgrade = True
|
||||
need_backup = True
|
||||
|
||||
if the_pkg_info is None:
|
||||
the_pkg_info = self.owner_env.lookup(pkg_id,version_str,cid)
|
||||
|
||||
if the_pkg_info is None:
|
||||
error_str = f"{pkg_name} ,cid:{cid} not found in index db"
|
||||
logger.error(error_str)
|
||||
return None,error_str
|
||||
|
||||
result_task = self.all_tasks.get(the_pkg_info.cid)
|
||||
if result_task is not None:
|
||||
return result_task,"already installing"
|
||||
|
||||
logger.info(f"start download&install {pkg_name},install_from_dependency={install_from_dependency},upgrade={is_upgrade},backup={need_backup},target_pkg_info={the_pkg_info}")
|
||||
result_task = PackageInstallTask(self.owner_env)
|
||||
self.all_tasks[the_pkg_info.cid] = result_task
|
||||
async def download_and_install_pkg()->int:
|
||||
# check dependency
|
||||
if skip_depends is False:
|
||||
result_task.dependency_tasks = {}
|
||||
self.get_dependency_tasks(the_pkg_info,result_task.dependency_tasks)
|
||||
result_task.state = INSTALL_TASK_STATE_INSTALL_DEPENDENCY
|
||||
for depend_pkg_name in result_task.dependency_tasks:
|
||||
# check pkg in local?
|
||||
# install miss pkg
|
||||
pass
|
||||
|
||||
result_task.state = INSTALL_TASK_STATE_DOWNLOADING
|
||||
install_full_path = ""
|
||||
target_full_path = ""
|
||||
old_package_full_path = ""
|
||||
is_download_directy = False
|
||||
|
||||
if the_pkg_info.target_media_type == the_pkg_info.source_media_type:
|
||||
is_download_directy = True
|
||||
if is_upgrade:
|
||||
target_full_path = ""
|
||||
else:
|
||||
target_full_path = ""
|
||||
else:
|
||||
pass
|
||||
|
||||
urls = self.owner_env.get_pkg_urls(the_pkg_info)
|
||||
#download
|
||||
client = NDN_Client() # set watch
|
||||
download_result = await client.get_file(the_pkg_info.cid,urls,target_full_path,options)
|
||||
if download_result !=0:
|
||||
result_task.state = INSTALL_TAKS_STATE_ERROR
|
||||
return result_task.state
|
||||
|
||||
result_task.state = INSTALL_TASK_STATE_INSTALLING
|
||||
if is_download_directy is False:
|
||||
install_media_result = False
|
||||
install_media_result = await self.owner_env.do_pkg_media_trans(the_pkg_info,target_full_path,install_full_path)
|
||||
if install_media_result is False:
|
||||
result_task.state = INSTALL_TAKS_STATE_ERROR
|
||||
result_task.error_str = "install media error,from {target_full_path} to {install_full_path}"
|
||||
return result_task.state
|
||||
|
||||
# last step,save install flag : install by manual or install by dependency
|
||||
## save cid dir
|
||||
if is_upgrade:
|
||||
os.rename(old_package_full_path, old_package_full_path + ".old" )
|
||||
os.rename(target_full_path,install_full_path)
|
||||
## update/create version link
|
||||
|
||||
## update pkg state
|
||||
## remove old version
|
||||
|
||||
result_task.state = INSTALL_TASK_STATE_DONE
|
||||
return result_task.state
|
||||
|
||||
|
||||
result_task.working_task = asyncio.create_task(download_and_install_pkg())
|
||||
return result_task,None
|
||||
|
||||
|
||||
def uninstall(self):
|
||||
pass
|
||||
|
||||
def get_dependency_tasks(self,pkg:PackageInfo,dependency_tasks):
|
||||
pass
|
||||
|
||||
async def check_dependency(self,pkg:PackageInfo,task_list:{}) -> bool:
|
||||
for depend_pkg_name in pkg.depends:
|
||||
depend_task = task_list.get(depend_pkg_name)
|
||||
if depend_task is not None:
|
||||
logger.debug(f"{pkg.name}'s depend pkg {depend_pkg_name} already in task list")
|
||||
continue
|
||||
depend_task = PackageInstallTask(self.owner_env)
|
||||
task_list[depend_pkg_name] = depend_task
|
||||
|
||||
depend_pkg_info = self.owner_env.lookup(depend_pkg_name)
|
||||
if depend_pkg_info is None:
|
||||
logger.warn(f"{pkg.name}'s depend pkg {depend_pkg_name} not found in index db")
|
||||
return False
|
||||
|
||||
if await self.check_dependency(depend_pkg_info,task_list) is False:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,18 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import aiofiles
|
||||
|
||||
class MediaReader(ABC):
|
||||
@abstractmethod
|
||||
async def read(self, inner_path:str,mode:str):
|
||||
pass
|
||||
|
||||
|
||||
class FolderMediaReader(MediaReader):
|
||||
def __init__(self, root_dir:str) -> None:
|
||||
self.root_dir = root_dir
|
||||
pass
|
||||
|
||||
async def read(self, inner_path:str,mode:str):
|
||||
full_path = self.root_dir + "/" + inner_path
|
||||
result_file = await aiofiles.open(full_path, mode,encoding='utf-8')
|
||||
return result_file
|
||||
@@ -1,41 +0,0 @@
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
class PackageInfo:
|
||||
def __init__(self) -> None:
|
||||
self.name = ""
|
||||
self.cid = None
|
||||
self.depends : list[str] = None
|
||||
self.author = None
|
||||
self.remote_urls = None
|
||||
self.target_media_type = "dir"
|
||||
self.source_media_type = "7z"
|
||||
|
||||
@staticmethod
|
||||
def parse_pkg_name(pkg_name:str) -> Tuple[str, str, str]:
|
||||
"""parse pkg name like test-pkg#nightly~>0.2.31#sha1:323423423 to test-pkg,nightly#>0.2.31,sha1:323423423"""
|
||||
args = pkg_name.split("#")
|
||||
if len(args) == 1:
|
||||
return args[0],None,None
|
||||
elif len(args) == 2:
|
||||
return args[0],None,arg[2]
|
||||
elif len(args) == 3:
|
||||
return args[0],args[1],args[2]
|
||||
else:
|
||||
logger.error(f"parse pkg name {pkg_name} failed!")
|
||||
return None,None,None
|
||||
|
||||
|
||||
|
||||
|
||||
@property
|
||||
def cid(self) -> str:
|
||||
return self.cid
|
||||
|
||||
class PackageMediaInfo:
|
||||
def __init__(self,full_path,media_type) -> None:
|
||||
self.media_type = media_type
|
||||
self.full_path = full_path
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
from .local_stability_node import *
|
||||
from .stability_node import *
|
||||
@@ -0,0 +1,203 @@
|
||||
import os
|
||||
import io
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
import base64
|
||||
from PIL import Image
|
||||
import requests
|
||||
from typing import Tuple
|
||||
from pathlib import Path
|
||||
|
||||
from aios import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeTaskResultCode,ComputeNode,AIStorage,UserConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Local_Stability_ComputeNode(ComputeNode):
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = Local_Stability_ComputeNode()
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
def declare_user_config(cls):
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
if os.getenv("LOCAL_STABILITY_URL") is None:
|
||||
user_config.add_user_config(
|
||||
"local_stability_url", "local stability url", True, None)
|
||||
if os.getenv("TEXT2IMG_OUTPUT_DIR") is None:
|
||||
home_dir = Path.home()
|
||||
output_dir = Path.joinpath(home_dir, "text2img_output")
|
||||
Path.mkdir(output_dir, exist_ok=True)
|
||||
user_config.add_user_config(
|
||||
"text2img_output_dir", "text2image output dir", True, output_dir)
|
||||
if os.getenv("TEXT2IMG_DEFAULT_MODEL") is None:
|
||||
user_config.add_user_config(
|
||||
"text2img_default_model", "text2img default model", True, "v1-5-pruned-emaonly")
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.is_start = False
|
||||
self.node_id = "local_stability_node"
|
||||
self.url = None
|
||||
self.default_model = None
|
||||
self.output_dir = None
|
||||
|
||||
self.task_queue = Queue()
|
||||
|
||||
async def initial(self):
|
||||
if os.getenv("LOCAL_STABILITY_URL") is not None:
|
||||
self.url = os.getenv("LOCAL_STABILITY_URL")
|
||||
else:
|
||||
self.url = AIStorage.get_instance(
|
||||
).get_user_config().get_value("local_stability_url")
|
||||
|
||||
if os.getenv("TEXT2IMG_OUTPUT_DIR") is not None:
|
||||
self.output_dir = os.getenv("TEXT2IMG_OUTPUT_DIR")
|
||||
else:
|
||||
self.output_dir = AIStorage.get_instance(
|
||||
).get_user_config().get_value("text2img_output_dir")
|
||||
|
||||
if os.getenv("TEXT2IMG_DEFAULT_MODEL") is not None:
|
||||
self.default_model = os.getenv("TEXT2IMG_DEFAULT_MODEL")
|
||||
else:
|
||||
self.default_model = AIStorage.get_instance(
|
||||
).get_user_config().get_value("text2img_default_model")
|
||||
|
||||
if self.url is None:
|
||||
logger.error("local stability url is None!")
|
||||
return False
|
||||
|
||||
if self.default_model is None:
|
||||
logger.error("local stability default model is None!")
|
||||
return False
|
||||
|
||||
if self.output_dir is None:
|
||||
self.output_dir = "./"
|
||||
|
||||
self.output_dir = os.path.abspath(self.output_dir)
|
||||
|
||||
self.start()
|
||||
|
||||
return True
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"stability_node push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
def _make_post_request(self, url, json) -> Tuple[str, requests.Response]:
|
||||
try:
|
||||
response = requests.post(url, json=json)
|
||||
if response.status_code != 200:
|
||||
return f'{response.status_code}, {response.json()}', None
|
||||
return None, response
|
||||
except Exception as e:
|
||||
return f"{e}", None
|
||||
|
||||
|
||||
def _run_task(self, task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
|
||||
model_name = task.params["model_name"]
|
||||
prompt = task.params["prompt"]
|
||||
negative_prompt = task.params["negative_prompt"]
|
||||
if negative_prompt == None or negative_prompt == "":
|
||||
negative_prompt = "sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, duplicate, mutated hands, mutated legs, (blurry:1.3), (bad anatomy:1.2), bad proportions, extra limbs, more than 2 nipples, extra legs, fused fingers, missing fingers, jpeg artifacts, signature, watermark, username, artist name, heterochromia, muscular legs, monochrome, grayscale, skin spots, acnes, skin blemishes, age spot, skin spots, acnes, logo, badhandv4, easynegative, cropped image, patreon,lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, ng_deepnegative_v1_75t, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry,(Tiptoe:1.3),looking at viewer, Twisted eyes"
|
||||
|
||||
prompt += ",masterpiece, best quality:1.3"
|
||||
|
||||
logging.info(f"call local stability {model_name} prompts: {prompt}, nagative_prompt: {negative_prompt}")
|
||||
|
||||
if model_name is not None:
|
||||
payload = {
|
||||
"sd_model_checkpoint": model_name,
|
||||
}
|
||||
err, resp = self._make_post_request(f'{self.url}/sdapi/v1/options', payload)
|
||||
|
||||
if err is not None:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
err_msg = f"Set local stability model failed. err:{err}"
|
||||
logger.error(err_msg)
|
||||
task.error_str = err_msg
|
||||
result.error_str = err_msg
|
||||
return result
|
||||
|
||||
logging.info(f"set local stability model {model_name} success")
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"steps": 20
|
||||
}
|
||||
|
||||
err, resp = self._make_post_request(f'{self.url}/sdapi/v1/txt2img', payload)
|
||||
if err is not None:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
err_msg = f"Failed. err:{err}"
|
||||
logger.error(err_msg)
|
||||
task.error_str = err_msg
|
||||
result.error_str = err_msg
|
||||
return result
|
||||
|
||||
r = resp.json()
|
||||
|
||||
for i in r['images']:
|
||||
image = Image.open(io.BytesIO(
|
||||
base64.b64decode(i.split(",", 1)[0])))
|
||||
file_name = os.path.join(self.output_dir, task.task_id + ".png")
|
||||
image.save(file_name)
|
||||
|
||||
task.state = ComputeTaskState.DONE
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.worker_id = self.node_id
|
||||
result.result = {"file": file_name}
|
||||
|
||||
return result
|
||||
|
||||
task.error_str = "Unknown error!"
|
||||
result.error_str = "Unknown error!"
|
||||
task.state = ComputeTaskState.ERROR
|
||||
return result
|
||||
|
||||
def start(self):
|
||||
if self.is_start:
|
||||
return
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
logger.info("local_stability_node is waiting for task...")
|
||||
task = await self.task_queue.get()
|
||||
logger.info(f"stability_node get task: {task.display()}")
|
||||
result = self._run_task(task)
|
||||
# if result is not None:
|
||||
# task.state = ComputeTaskState.DONE
|
||||
# task.result = result
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
def display(self) -> str:
|
||||
return f"Stability_ComputeNode: {self.node_id}"
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
return task.task_type == ComputeTaskType.TEXT_2_IMAGE
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return False
|
||||
@@ -0,0 +1,199 @@
|
||||
import os
|
||||
import io
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
from stability_sdk import client
|
||||
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
|
||||
|
||||
from aios import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeTaskResultCode,ComputeNode,AIStorage,UserConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Stability_ComputeNode(ComputeNode):
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = Stability_ComputeNode()
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
def declare_user_config(cls):
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
user_config.add_user_config(
|
||||
"stability_api_key", "stability api key", False, None)
|
||||
user_config.add_user_config(
|
||||
"stability_model", "stability model name", True, "stable-diffusion-512-v2-1")
|
||||
if os.getenv("TEXT2IMG_OUTPUT_DIR") is None:
|
||||
home_dir = Path.home()
|
||||
output_dir = Path.joinpath(home_dir, "text2img_output")
|
||||
Path.mkdir(output_dir, exist_ok=True)
|
||||
user_config.add_user_config(
|
||||
"text2img_output_dir", "text2image output dir", True, output_dir)
|
||||
if os.getenv("STABILITY_DEFAULT_MODEL") is None:
|
||||
user_config.add_user_config(
|
||||
"stability_default_model", "stability default model", True, "stable-diffusion-512-v2-1")
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.is_start = False
|
||||
self.node_id = "stability_node"
|
||||
self.api_key = ""
|
||||
self.default_model = ""
|
||||
|
||||
self.task_queue = Queue()
|
||||
|
||||
async def initial(self):
|
||||
if os.getenv("STABILITY_API_KEY") is not None:
|
||||
self.api_key = os.getenv("STABILITY_API_KEY")
|
||||
else:
|
||||
self.api_key = AIStorage.get_instance(
|
||||
).get_user_config().get_value("stability_api_key")
|
||||
|
||||
if self.api_key is None:
|
||||
logger.error("stability api key is None!")
|
||||
return False
|
||||
|
||||
# Check out the following link for a list of available engines: https://platform.stability.ai/docs/features/api-parameters#engine
|
||||
if os.getenv("STABILITY_DEFAULT_MODEL") is not None:
|
||||
self.default_model = os.getenv("STABILITY_DEFAULT_MODEL")
|
||||
else:
|
||||
self.default_model = AIStorage.get_instance().get_user_config().get_value("stability_default_model")
|
||||
|
||||
if self.default_model is None:
|
||||
self.default_model = "stable-diffusion-512-v2-1"
|
||||
|
||||
if os.getenv("TEXT2IMG_OUTPUT_DIR") is not None:
|
||||
self.output_dir = os.getenv("TEXT2IMG_OUTPUT_DIR")
|
||||
else:
|
||||
self.output_dir = AIStorage.get_instance(
|
||||
).get_user_config().get_value("text2img_output_dir")
|
||||
|
||||
if self.output_dir is None:
|
||||
self.output_dir = "./"
|
||||
self.output_dir = os.path.abspath(self.output_dir)
|
||||
|
||||
self.start()
|
||||
|
||||
return True
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"stability_node push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
def _run_task(self, task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
|
||||
model_name = task.params["model_name"]
|
||||
prompt = task.params["prompt"]
|
||||
negative_prompt = task.params["negative_prompt"]
|
||||
|
||||
logging.info(f"call stability {self.default_model} prompts: {prompt}, negative_prompt: {negative_prompt}")
|
||||
|
||||
api = None
|
||||
try:
|
||||
api = client.StabilityInference(
|
||||
key=self.api_key,
|
||||
verbose=True, # Print debug messages.
|
||||
engine=model_name,
|
||||
)
|
||||
except Exception as e:
|
||||
task.error_str = f"create stability client failed: {e}"
|
||||
result.error_str = f"create stability client failed: {e}"
|
||||
logging.warn(task.error_str)
|
||||
task.state = ComputeTaskState.ERROR
|
||||
return result
|
||||
|
||||
answers = api.generate(
|
||||
prompt=prompt,
|
||||
# If a seed is provided, the resulting generated image will be deterministic.
|
||||
seed=0,
|
||||
# What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again.
|
||||
# Note: This isn't quite the case for Clip Guided generations, which we'll tackle in a future example notebook.
|
||||
# Amount of inference steps performed on image generation. Defaults to 30.
|
||||
steps=30,
|
||||
# Influences how strongly your generation is guided to match your prompt.
|
||||
cfg_scale=7.0,
|
||||
# Setting this value higher increases the strength in which it tries to match your prompt.
|
||||
# Defaults to 7.0 if not specified.
|
||||
width=512, # Generation width, defaults to 512 if not included.
|
||||
height=512, # Generation height, defaults to 512 if not included.
|
||||
# Number of images to generate, defaults to 1 if not included.
|
||||
samples=1,
|
||||
# Choose which sampler we want to denoise our generation with.
|
||||
sampler=generation.SAMPLER_K_DPMPP_2M
|
||||
# Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers.
|
||||
# (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m, k_dpmpp_sde)
|
||||
)
|
||||
|
||||
for resp in answers:
|
||||
for artifact in resp.artifacts:
|
||||
if artifact.finish_reason == generation.FILTER:
|
||||
err_msg = "request activated the API's safety filters"
|
||||
logging.warn(err_msg)
|
||||
task.error_str = err_msg
|
||||
result.error_str = err_msg
|
||||
task.state = ComputeTaskState.ERROR
|
||||
return result
|
||||
if artifact.type == generation.ARTIFACT_IMAGE:
|
||||
img = Image.open(io.BytesIO(artifact.binary))
|
||||
# Save our generated images with the task_id as the filename.
|
||||
file_name = os.path.join(self.output_dir, task.task_id + ".png")
|
||||
img.save(file_name)
|
||||
|
||||
task.state = ComputeTaskState.DONE
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.worker_id = self.node_id
|
||||
result.result = {"file": file_name}
|
||||
|
||||
return result
|
||||
|
||||
task.error_str = "Unknown error!"
|
||||
result.error_str = "Unknown error!"
|
||||
task.state = ComputeTaskState.ERROR
|
||||
return result
|
||||
|
||||
def start(self):
|
||||
if self.is_start:
|
||||
return
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
logger.info("stability_node is waiting for task...")
|
||||
task = await self.task_queue.get()
|
||||
logger.info(f"stability_node get task: {task.display()}")
|
||||
result = self._run_task(task)
|
||||
# if result is not None:
|
||||
# task.state = ComputeTaskState.DONE
|
||||
# task.result = result
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
def display(self) -> str:
|
||||
return f"Stability_ComputeNode: {self.node_id}"
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
return task.task_type == ComputeTaskType.TEXT_2_IMAGE
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return False
|
||||
@@ -0,0 +1 @@
|
||||
from .local_st_compute_node import *
|
||||
@@ -0,0 +1,212 @@
|
||||
import logging
|
||||
import requests
|
||||
from typing import Optional, List
|
||||
from pydantic import BaseModel
|
||||
from typing import Union
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
from aios import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeTaskResultCode,ComputeNode,AIStorage,UserConfig,ObjectID,Queue_ComputeNode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
|
||||
# For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html
|
||||
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
|
||||
super().__init__()
|
||||
|
||||
self.node_id = "local_sentence_transformer_text_embedding_node"
|
||||
self.model_name = model_name
|
||||
self.model = None
|
||||
|
||||
def initial(self) -> bool:
|
||||
logger.info(
|
||||
f"LocalSentenceTransformer_Text_ComputeNode init, model_name: {self.model_name}"
|
||||
)
|
||||
|
||||
assert self.model_name is not None
|
||||
assert self.model is None
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
self.model = SentenceTransformer(self.model_name)
|
||||
except Exception as err:
|
||||
logger.error(f"load model {self.model} failed: {err}")
|
||||
return False
|
||||
self.start()
|
||||
return True
|
||||
|
||||
async def execute_task(self, task: ComputeTask) :
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
try:
|
||||
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}")
|
||||
if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
|
||||
input = task.params["input"]
|
||||
logger.debug(
|
||||
f"LocalSentenceTransformer_Text_ComputeNode task input: {input}"
|
||||
)
|
||||
sentence_embeddings = self.model.encode(input, show_progress_bar=False).tolist()
|
||||
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.result["content"] = sentence_embeddings
|
||||
|
||||
else:
|
||||
result.error_str = f"unsupport embedding task type: {task.task_type}"
|
||||
except Exception as err:
|
||||
import traceback
|
||||
|
||||
logger.error(f"{traceback.format_exc()}, error: {err}")
|
||||
result.error_str = f"{traceback.format_exc()}, error: {err}"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def display(self) -> str:
|
||||
return f"LocalSentenceTransformer_Text_ComputeNode: {self.node_id}, {self.model_name}"
|
||||
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
return task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "all-MiniLM-L6-v2"
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
|
||||
# For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "clip-ViT-B-32",
|
||||
multi_model_name: str = "clip-ViT-B-32-multilingual-v1",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.node_id = "local_sentence_transformer_image_embedding_node"
|
||||
self.model_name = model_name
|
||||
self.multi_model_name = multi_model_name
|
||||
self.model = None
|
||||
self.multi_model = None
|
||||
|
||||
def initial(self) -> bool:
|
||||
logger.info(
|
||||
f"LocalSentenceTransformer_Image_ComputeNode init, model_name: {self.model_name} {self.multi_model_name}"
|
||||
)
|
||||
|
||||
assert self.model_name is not None
|
||||
assert self.multi_model_name is not None
|
||||
assert self.model is None
|
||||
assert self.multi_model is None
|
||||
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
self.model = SentenceTransformer(self.model_name)
|
||||
self.multi_model = SentenceTransformer(self.multi_model_name)
|
||||
except Exception as err:
|
||||
logger.error(f"load model {self.model} failed: {err}")
|
||||
return False
|
||||
self.start()
|
||||
return True
|
||||
|
||||
def _load_image(self, source: Union[ObjectID, bytes]) -> Optional[Image]:
|
||||
image_data = None
|
||||
if isinstance(source, ObjectID):
|
||||
from knowledge import KnowledgeStore, ImageObject
|
||||
|
||||
buf = KnowledgeStore().get_object_store().get_object(source)
|
||||
if buf is None:
|
||||
logger.error(f"load image object but not found! {source}")
|
||||
return None
|
||||
|
||||
try:
|
||||
image_obj = ImageObject.decode(buf)
|
||||
except Exception as err:
|
||||
logger.error(f"decode ImageObject from buffer failed: {source}, {err}")
|
||||
return None
|
||||
|
||||
file_size = image_obj.get_file_size()
|
||||
# print(f"got image object: {source.to_base58()}, size: {file_size}")
|
||||
|
||||
image_data = (
|
||||
KnowledgeStore()
|
||||
.get_chunk_reader()
|
||||
.read_chunk_list_to_single_bytes(image_obj.get_chunk_list())
|
||||
)
|
||||
|
||||
elif isinstance(source, bytes):
|
||||
image_data = source
|
||||
else:
|
||||
logger.error(f"unsupport image source type: {type(source)}, {source}")
|
||||
return None
|
||||
|
||||
try:
|
||||
img = Image.open(io.BytesIO(image_data))
|
||||
|
||||
return img
|
||||
except Exception as err:
|
||||
logger.error(f"load image from buffer failed: {source}, {err}")
|
||||
return None
|
||||
|
||||
async def execute_task(
|
||||
self, task: ComputeTask
|
||||
) -> ComputeTaskResult:
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
try:
|
||||
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}")
|
||||
if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
|
||||
input = task.params["input"]
|
||||
logger.debug(
|
||||
f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}"
|
||||
)
|
||||
sentence_embeddings = self.multi_model.encode(input, show_progress_bar=False).tolist()
|
||||
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.result["content"] = sentence_embeddings
|
||||
|
||||
elif task.task_type == ComputeTaskType.IMAGE_EMBEDDING:
|
||||
input = task.params["input"]
|
||||
logger.debug(
|
||||
f"LocalSentenceTransformer_Image_ComputeNode task image input: {input}"
|
||||
)
|
||||
|
||||
img = self._load_image(input)
|
||||
if img is None:
|
||||
result.error_str = f"load image failed: {input}"
|
||||
return result
|
||||
|
||||
sentence_embeddings = self.model.encode(img, show_progress_bar=False).tolist()
|
||||
result.result_code = ComputeTaskResultCode.OK
|
||||
result.result["content"] = sentence_embeddings
|
||||
else:
|
||||
result.error_str = f"unsupport embedding task type: {task.task_type}"
|
||||
except Exception as err:
|
||||
import traceback
|
||||
|
||||
logger.error(f"{traceback.format_exc()}, error: {err}")
|
||||
result.error_str = f"{traceback.format_exc()}, error: {err}"
|
||||
|
||||
|
||||
return result
|
||||
|
||||
def display(self) -> str:
|
||||
return f"LocalSentenceTransformer_Image_ComputeNode: {self.node_id}, {self.model_name}"
|
||||
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
return (
|
||||
(task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "clip-ViT-B-32")
|
||||
or task.task_type == ComputeTaskType.IMAGE_EMBEDDING
|
||||
)
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return True
|
||||
@@ -0,0 +1,307 @@
|
||||
import logging
|
||||
import threading
|
||||
import asyncio
|
||||
import uuid
|
||||
import time
|
||||
import aiofiles
|
||||
|
||||
from telegram import Update,Message
|
||||
from telegram import Bot
|
||||
from telegram.ext import Updater
|
||||
from telegram.error import Forbidden, NetworkError
|
||||
|
||||
from aios import ObjectType, KnowledgeStore,AgentTunnel,AIStorage,ContactManager,Contact,FamilyMember,AgentMsg,AgentMsgType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TelegramTunnel(AgentTunnel):
|
||||
all_bots = {}
|
||||
default_chatid = {}
|
||||
@classmethod
|
||||
def register_to_loader(cls):
|
||||
async def load_tg_tunnel(config:dict) -> AgentTunnel:
|
||||
result_tunnel = TelegramTunnel("")
|
||||
if await result_tunnel.load_from_config(config):
|
||||
return result_tunnel
|
||||
else:
|
||||
return None
|
||||
|
||||
AgentTunnel.register_loader("TelegramTunnel",load_tg_tunnel)
|
||||
|
||||
|
||||
async def load_from_config(self,config:dict)->bool:
|
||||
self.type = "TelegramTunnel"
|
||||
self.tg_token = config["token"]
|
||||
self.target_id = config["target"]
|
||||
self.tunnel_id = config["tunnel_id"]
|
||||
if config.get("allow") is not None:
|
||||
self.allow_group = config["allow"]
|
||||
|
||||
return True
|
||||
|
||||
def dump_to_config(self) -> dict:
|
||||
pass
|
||||
|
||||
def __init__(self,tg_token:str) -> None:
|
||||
super().__init__()
|
||||
self.is_start = False
|
||||
self.tg_token = tg_token
|
||||
self.bot:Bot = None
|
||||
self.update_queue = None
|
||||
self.allow_group = "contact"
|
||||
self.in_process_tg_msg = {}
|
||||
self.chatid_record = {}
|
||||
|
||||
async def _do_process_raw_message(self,bot: Bot, update_id: int) -> int:
|
||||
# Request updates after the last update_id
|
||||
updates = await bot.get_updates(offset=update_id, timeout=10, allowed_updates=Update.ALL_TYPES)
|
||||
for update in updates:
|
||||
next_update_id = update.update_id + 1
|
||||
|
||||
if update.message and update.message.text:
|
||||
|
||||
await self.on_message(bot,update)
|
||||
return next_update_id
|
||||
|
||||
return update_id
|
||||
|
||||
async def start(self) -> bool:
|
||||
if self.is_start:
|
||||
logger.warning(f"tunnel {self.tunnel_id} is already started")
|
||||
return False
|
||||
self.is_start = True
|
||||
logger.info(f"tunnel {self.tunnel_id} is starting...")
|
||||
|
||||
self.bot = Bot(self.tg_token)
|
||||
self.bot_username = (await self.bot.get_me()).username
|
||||
self.update_queue = asyncio.Queue()
|
||||
self.bot_updater = Updater(self.bot,update_queue=self.update_queue)
|
||||
|
||||
TelegramTunnel.all_bots[self.target_id] = self.bot
|
||||
|
||||
async def _run_app():
|
||||
try:
|
||||
update_id = (await self.bot.get_updates())[0].update_id
|
||||
except IndexError:
|
||||
update_id = None
|
||||
|
||||
#logger.info("listening for new messages...")
|
||||
while True:
|
||||
try:
|
||||
update_id = await self._do_process_raw_message(self.bot, update_id)
|
||||
except NetworkError:
|
||||
await asyncio.sleep(1)
|
||||
except Forbidden:
|
||||
# The user has removed or blocked the bot.
|
||||
update_id += 1
|
||||
except Exception as e:
|
||||
logger.error(f"tg_tunnel error:{e}")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
|
||||
asyncio.create_task(_run_app())
|
||||
logger.info(f"tunnel {self.tunnel_id} started.")
|
||||
return True
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
async def _process_message(self, msg: AgentMsg) -> bool:
|
||||
logger.warn(f"tg_tunnel process message {msg.msg_id} from agent {msg.sender} to human {msg.target}")
|
||||
|
||||
# async def _process_message(self, msg: AgentMsg) -> bool:
|
||||
# logger.info(f"tg_tunnel process message {msg.msg_id} from agent {msg.sender} to human {msg.target}")
|
||||
# cm = ContactManager.get_instance()
|
||||
# contact = cm.find_contact_by_name(msg.target)
|
||||
# bot = TelegramTunnel.all_bots.get(msg.sender)
|
||||
# chatid_index = f"{self.target_id}#{msg.target}"
|
||||
# chatid = TelegramTunnel.default_chatid.get(chatid_index)
|
||||
# if chatid is None:
|
||||
# logger.warning(f"tg_tunnel process message {msg.msg_id} from agent {msg.sender} to human {msg.target} failed! chatid not found!")
|
||||
# return None
|
||||
|
||||
# if bot is None:
|
||||
# logger.warning(f"tg_tunnel process message {msg.msg_id} from agent {msg.sender} to human {msg.target} failed! bot not found!")
|
||||
# return None
|
||||
|
||||
# if contact:
|
||||
# if contact.telegram:
|
||||
# await bot.send_message(chat_id=chatid,text=msg.body)
|
||||
# logging.info(f"tg_tunnel send message {msg.msg_id} from agent {msg.sender} to human {msg.target} @ chatid:{chatid}success!")
|
||||
# return None
|
||||
|
||||
# logger.warning(f"tg_tunnel process message {msg.msg_id} from agent {msg.sender} to human {msg.target} failed! contact not found!")
|
||||
# return None
|
||||
|
||||
async def post_message(self, msg: AgentMsg) -> None:
|
||||
chatid = self.chatid_record.get(msg.target)
|
||||
if chatid:
|
||||
# TODO: support image and audio
|
||||
await self.bot.send_message(chat_id=chatid,text=msg.body)
|
||||
logging.info(f"tg_tunnel send message {msg.msg_id} from agent {msg.sender} to human {msg.target} @ chatid:{chatid}success!")
|
||||
else:
|
||||
logger.warning(f"tg_tunnel process message {msg.msg_id} from agent {msg.sender} to human {msg.target} failed! chatid not found!")
|
||||
|
||||
|
||||
async def conver_tg_msg_to_agent_msg(self,message:Message) -> AgentMsg:
|
||||
agent_msg = AgentMsg()
|
||||
agent_msg.topic = "_telegram"
|
||||
agent_msg.msg_id = "tg_msg#" + str(message.message_id) + "#" + uuid.uuid4().hex
|
||||
agent_msg.target = self.target_id
|
||||
agent_msg.body = message.text
|
||||
agent_msg.create_time = time.time()
|
||||
messag_type = message.chat.type
|
||||
if messag_type == "supergroup" or messag_type == "group":
|
||||
agent_msg.target = f"tg_group{message.chat_id}"
|
||||
agent_msg.msg_type = AgentMsgType.TYPE_GROUPMSG
|
||||
agent_msg.mentions = []
|
||||
else:
|
||||
agent_msg.msg_type = AgentMsgType.TYPE_MSG
|
||||
agent_msg.mentions = []
|
||||
|
||||
if message.entities:
|
||||
for entity in message.entities:
|
||||
if entity.type == 'mention':
|
||||
mention = message.text[entity.offset:entity.offset+entity.length]
|
||||
if mention == '@' + self.bot_username:
|
||||
agent_msg.mentions.append(self.target_id)
|
||||
else:
|
||||
agent_msg.mentions.append(mention)
|
||||
|
||||
if message.caption_entities:
|
||||
for entity in message.caption_entities:
|
||||
if entity.type == 'mention':
|
||||
mention = message.caption[entity.offset:entity.offset+entity.length]
|
||||
if mention == '@' + self.bot_username:
|
||||
agent_msg.mentions.append(self.target_id)
|
||||
else:
|
||||
agent_msg.mentions.append(mention)
|
||||
|
||||
return agent_msg
|
||||
|
||||
def is_bot_mentioned(self,message:Message):
|
||||
if message.entities:
|
||||
for entity in message.entities:
|
||||
if entity.type == 'mention':
|
||||
mention = message.text[entity.offset:entity.offset+entity.length]
|
||||
if mention == '@' + self.bot_username:
|
||||
return True
|
||||
|
||||
if message.caption_entities:
|
||||
for entity in message.caption_entities:
|
||||
if entity.type == 'mention':
|
||||
mention = message.caption[entity.offset:entity.offset+entity.length]
|
||||
if mention == '@' + self.bot_username:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def on_message(self, bot:Bot, update: Update) -> None:
|
||||
message = update.message
|
||||
logger.info(f"on_message: {message.message_id} from {message.from_user.username} ({update.effective_user.username}) to {message.chat.title}({message.chat.id})")
|
||||
if update.effective_user.is_bot:
|
||||
logger.warning(f"ignore message from telegram bot {update.effective_user.id}")
|
||||
return None
|
||||
|
||||
if self.in_process_tg_msg.get(update.message.message_id) is not None:
|
||||
logger.warning(f"ignore message from telegram bot {update.effective_user.id}")
|
||||
return None
|
||||
|
||||
self.in_process_tg_msg[update.message.message_id] = True
|
||||
|
||||
agent_msg = await self.conver_tg_msg_to_agent_msg(message)
|
||||
cm : ContactManager = ContactManager.get_instance()
|
||||
reomte_user_name = f"{update.effective_user.id}@telegram"
|
||||
|
||||
contact : Contact = cm.find_contact_by_telegram(update.effective_user.username)
|
||||
if contact is None:
|
||||
contact = cm.find_contact_by_telegram(str(update.effective_user.id))
|
||||
|
||||
if contact is not None:
|
||||
reomte_user_name = contact.name
|
||||
|
||||
#TelegramTunnel.default_chatid[f"{self.target_id}#{reomte_user_name}"] = update.effective_chat.id
|
||||
if not contact.is_family_member:
|
||||
if self.allow_group != "contact" and self.allow_group !="guest":
|
||||
await update.message.reply_text(f"You're not supposed to talk to me! Please contact my father~")
|
||||
return
|
||||
|
||||
else:
|
||||
if self.allow_group != "guest":
|
||||
await update.message.reply_text(f"The current Telegram account is not in the contact list. If you want to receive a reply, you can add the configuration in the contacts.toml file or switch tunnel to guest mode.")
|
||||
return
|
||||
|
||||
if cm.is_auto_create_contact_from_telegram:
|
||||
contact_name = update.effective_user.first_name
|
||||
if update.effective_user.last_name is not None:
|
||||
contact_name += " " + update.effective_user.last_name
|
||||
|
||||
contact = Contact(contact_name)
|
||||
contact.telegram = update.effective_user.username if update.effective_user.username is not None else str(update.effective_user.id)
|
||||
contact.added_by = self.target_id
|
||||
cm.add_contact(contact.name, contact)
|
||||
reomte_user_name = contact.name
|
||||
|
||||
if contact is not None:
|
||||
contact.set_active_tunnel(self.target_id,self)
|
||||
self.chatid_record[reomte_user_name] = update.effective_chat.id
|
||||
self.ai_bus.register_message_handler(reomte_user_name,contact._process_msg)
|
||||
|
||||
agent_msg.sender = reomte_user_name
|
||||
logger.info(f"process message {agent_msg.msg_id} from {agent_msg.sender} to {agent_msg.target}")
|
||||
if agent_msg.msg_type == AgentMsgType.TYPE_GROUPMSG:
|
||||
self.ai_bus.register_message_handler(agent_msg.target, self._process_message)
|
||||
resp_msg = await self.ai_bus.send_message(agent_msg,self.target_id,agent_msg.target)
|
||||
else:
|
||||
#self.ai_bus.register_message_handler(reomte_user_name, self._process_message)
|
||||
resp_msg = await self.ai_bus.send_message(agent_msg)
|
||||
#await bot.send_chat_action(chat_id=update.effective_chat.id, action="typing")
|
||||
|
||||
|
||||
|
||||
if resp_msg is None:
|
||||
await update.message.reply_text(f"System Error: Timeout,{self.target_id} no resopnse! Please check logs/aios.log for more details!")
|
||||
else:
|
||||
if resp_msg.body_mime is None:
|
||||
if resp_msg.body is None:
|
||||
return
|
||||
|
||||
if len(resp_msg.body) < 1:
|
||||
await update.message.reply_text("")
|
||||
return
|
||||
|
||||
knowledge_object = KnowledgeStore().parse_object_in_message(resp_msg.body)
|
||||
if knowledge_object is not None:
|
||||
if knowledge_object.get_object_type() == ObjectType.Image:
|
||||
image = KnowledgeStore().bytes_from_object(knowledge_object)
|
||||
try:
|
||||
async with aiofiles.open("tg_send_temp.png", mode='wb') as local_file:
|
||||
if local_file:
|
||||
await local_file.write(image)
|
||||
await update.message.reply_photo("tg_send_temp.png")
|
||||
except Exception as e:
|
||||
logger.error(f"save image error: {e}")
|
||||
return
|
||||
else:
|
||||
pos = resp_msg.body.find("audio file")
|
||||
if pos != -1:
|
||||
audio_file = resp_msg.body[pos+11:].strip()
|
||||
if audio_file.startswith("\""):
|
||||
audio_file = audio_file[1:-1]
|
||||
await update.message.reply_voice(audio_file)
|
||||
return
|
||||
await update.message.reply_text(resp_msg.body)
|
||||
else:
|
||||
if resp_msg.body_mime.startswith("image"):
|
||||
photo_file = open(resp_msg.body,"rb")
|
||||
if photo_file:
|
||||
await update.message.reply_photo(resp_msg.body)
|
||||
photo_file.close()
|
||||
else:
|
||||
await update.message.reply_text(resp_msg.body)
|
||||
|
||||
else:
|
||||
await update.message.reply_text(resp_msg.body)
|
||||
|
||||
|
||||
@@ -2,8 +2,7 @@ import logging
|
||||
import toml
|
||||
import os
|
||||
|
||||
from aios_kernel import Workflow,AIStorage
|
||||
from package_manager import PackageEnv,PackageEnvManager,PackageMediaInfo,PackageInstallTask
|
||||
from aios import Workflow,AIStorage,PackageEnv,PackageEnvManager,PackageMediaInfo,PackageInstallTask
|
||||
from agent_manager import AgentManager
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user