Refactor the code directory structure to better suit the current complexity

This commit is contained in:
Liu Zhicong
2023-11-30 21:04:19 -08:00
parent 4955225ecd
commit adeca91e0a
99 changed files with 391 additions and 342 deletions
+1 -3
View File
@@ -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__)
+174
View File
@@ -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}")
+1
View File
@@ -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)
+1 -1
View File
@@ -2,7 +2,7 @@ import os
import runpy
import toml
import asyncio
from knowledge import KnowledgePipelineEnvironment, KnowledgePipeline
from aios import KnowledgePipelineEnvironment, KnowledgePipeline
class KnowledgePipelineManager:
+1
View File
@@ -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 -2
View File
@@ -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):
+1 -2
View File
@@ -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
+1 -1
View File
@@ -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:
+1 -2
View File
@@ -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:
-2
View File
@@ -1,2 +0,0 @@
from .cid import ContentId
from .ndn_client import NDN_Client
-12
View File
@@ -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
-117
View File
@@ -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
+3
View File
@@ -0,0 +1,3 @@
from .open_ai_node import *
from .openai_tts_node import *
from .whisper_node import *
+301
View File
@@ -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
+226
View File
@@ -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
View File
@@ -1 +0,0 @@
TODO
@@ -1,3 +0,0 @@
from .env import PackageEnvManager,PackageEnv
from .pkg import PackageInfo,PackageMediaInfo
from .installer import PackageInstallTask
-158
View File
@@ -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 @@
-170
View File
@@ -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
-41
View 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
+2
View File
@@ -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
+199
View File
@@ -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
+1
View File
@@ -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
+307
View File
@@ -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__)