sigmaflow 0.0.1__py3-none-any.whl

sigmaflow/__init__.py

@@ -0,0 +1,8 @@
+from .prompt import Prompt
+from .pipeline import Pipeline
+from .manager import PromptManager, PipelineManager
+from .server import PipelineServer
+from .api import PipelineAPI
+from .log import log
+
+__version__ = "0.0.1"

sigmaflow/api.py

@@ -0,0 +1,86 @@
+import os
+from fastapi import APIRouter
+from pydantic import BaseModel
+from fastapi import HTTPException
+from . import PromptManager, PipelineManager
+
+class PromptData(BaseModel):
+    name:   str | None = None
+    text:   str | None = None
+    keys:  list | None = None
+
+class PipeData(BaseModel):
+    name:   str | None = None
+    data:  dict | None = None
+
+class PipelineAPI:
+    def __init__(self, prompt_dir, pipeline_dir, llm_client=None, rag_client=None, is_async=False):
+        router = APIRouter(prefix='/api')
+        prompt_manager = PromptManager(prompt_dir)
+        pipeline_manager = PipelineManager(llm_client, rag_client, is_async, pipeline_dir, prompt_manager)
+
+        @router.get("/prompt/list")
+        async def prompt_list():
+            try:
+                ret = {k: {'text': v.text, 'keys': v.keys} for k,v in prompt_manager.prompts.items()}
+                return ret
+
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        @router.post("/prompt/update")
+        async def prompt_update(p_data: PromptData):
+            try:
+                if p_data.text:
+                    prompt_manager.prompts[p_data.name].text = p_data.text
+                if p_data.keys:
+                    prompt_manager.prompts[p_data.name].keys = p_data.keys
+
+                ret = {
+                    p_data.name: {
+                        'text': prompt_manager.prompts[p_data.name].text,
+                        'keys': prompt_manager.prompts[p_data.name].keys
+                    }
+                }
+                return ret
+
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+
+        @router.get("/pipeline/list")
+        async def pipeline_list():
+            try:
+                ret = pipeline_manager.export_pipe_conf()
+                return ret
+
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        @router.post("/pipeline/mermaid")
+        async def pipeline_mermaid(p_data: PipeData):
+            try:
+                pipeline = pipeline_manager.pipes[p_data.name]['func']
+                mermaid = pipeline.pipe2mermaid(pipeline.pipe)
+                ret = {
+                    'mermaid': mermaid,
+                }
+                return ret
+
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        @router.post("/pipeline/update")
+        async def pipeline_update(p_data: PipeData):
+            try:
+                ret = {
+                    'result': pipeline_manager.update_pipe(p_data.name, p_data.data)
+                }
+                return ret
+
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        self.router = router
+        self.prompt_manager = prompt_manager
+        self.pipeline_manager = pipeline_manager

sigmaflow/clients/__init__.py

@@ -0,0 +1,22 @@
+import os
+import asyncio
+import pkgutil
+import importlib
+from ..log import log
+
+class _SearchEngine:
+    def __init__(self):
+        log.debug('Setup SearchEngine')
+        self.engines = {}
+        for loader, module_name, is_pkg in pkgutil.walk_packages(__path__, __name__ + '.'):
+            if (name := module_name.split('.')[-1]).startswith('search_'):
+                name = name.replace('search_', '')
+                m = importlib.import_module(module_name)
+                self.engines[name] = [m.search, m.async_search]
+        log.debug(f'Load engines: {list(self.engines.keys())}')
+
+    def get(self, engine_name):
+        return self.engines[engine_name]
+
+SearchEngine = _SearchEngine()
+llm_sem = asyncio.Semaphore(int(os.getenv('LLM_CONCURRENCY', 16)))

sigmaflow/clients/llm_lmdeploy.py

@@ -0,0 +1,87 @@
+from . import *
+spec = importlib.util.find_spec('lmdeploy')
+
+if spec:
+    from lmdeploy import pipeline, GenerationConfig, TurbomindEngineConfig, ChatTemplateConfig
+
+    model_path = os.getenv('LMDEPLOY_MODEL')
+    chat_temp_name = os.getenv('LMDEPLOY_CHAT_TEMP', None)
+    session_len = int(os.getenv('LMDEPLOY_SESSION_LEN', 256))
+    max_new_tokens = int(os.getenv('LMDEPLOY_MAX_NEW_TOKENS', 256))
+    batch_size = int(os.getenv('LMDEPLOY_BATCH_SIZE', 128))
+    pp = int(os.getenv('LMDEPLOY_PP', 1))
+    tp = int(os.getenv('LMDEPLOY_TP', 1))
+
+    backend_config = TurbomindEngineConfig(
+                        session_len=session_len,
+                        pp=pp,
+                        tp=tp)
+    gen_config = GenerationConfig(
+                        do_sample=True,
+                        # top_p=0.95,
+                        # temperature=0.7,
+                        max_new_tokens=max_new_tokens)
+    chat_temp = ChatTemplateConfig(model_name=chat_temp_name) if chat_temp_name else None
+    llm = pipeline(model_path, backend_config=backend_config, chat_template_config=chat_temp)
+
+batch_wait_time = 1
+batch_queue = []
+batch_event = asyncio.Event()
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    raise NotImplementedError
+
+def batch_completion(batch_data):
+    prompts = []
+    for text, _ in batch_data:
+        if (t := type(text)) is str:
+            prompts.append(text)
+        elif t is list:
+            msg = []
+            for role, c in zip(['user', 'assistant']*len(text), text):
+                msg.append({
+                    "role": role,
+                    "content": c,
+                })
+            prompts.append(msg)
+
+    outputs = llm(prompts, gen_config=gen_config)
+    result = [o.text for o in outputs]
+    return result
+
+async def llm_batch_processor():
+    global batch_queue
+    while True:
+        await batch_event.wait()
+
+        if batch_queue:
+            batch = batch_queue[:batch_size]
+            batch_queue = batch_queue[len(batch):]
+            result = await asyncio.to_thread(batch_completion, batch)
+
+            for (_, future), res in zip(batch, result):
+                future.set_result(res)
+
+        if not batch_queue:
+            batch_event.clear()
+
+async def async_completion(text):
+    global batch_queue
+
+    loop = asyncio.get_event_loop()
+    future = loop.create_future()
+    batch_queue.append((text, future))
+
+    if len(batch_queue) >= batch_size:
+        batch_event.set()
+    else:
+        await asyncio.sleep(batch_wait_time)
+        batch_event.set()
+
+    return await future

sigmaflow/clients/llm_mlx.py

@@ -0,0 +1,24 @@
+import os
+from mlx_lm import load, generate
+from . import *
+
+model, tokenizer = load(os.getenv('MLX_MODEL'))
+max_tokens = int(os.getenv('MLX_MAX_TOKENS', 256))
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    messages = [{"role": "user", "content": text}]
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    text = generate(model, tokenizer, prompt=prompt, verbose=False, max_tokens=max_tokens)
+
+    return text
+
+async def async_completion(text):
+    return completion(text)

sigmaflow/clients/llm_mlx_batch.py

@@ -0,0 +1,63 @@
+import os, sys
+sys.path.append('/Users/mkk/workspace/git_repos/mlx_parallm')
+from mlx_parallm.utils import load, batch_generate
+from . import *
+
+model, tokenizer = load(os.getenv('MLX_MODEL'))
+max_tokens = int(os.getenv('MLX_MAX_TOKENS', 256))
+batch_size = int(os.getenv('MLX_BATCH_SIZE', 128))
+
+batch_wait_time = 1
+batch_queue = []
+batch_event = asyncio.Event()
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    raise NotImplementedError
+
+def batch_completion(batch_data):
+    print(f"Processed batch, len: {len(batch_data)}")
+    prompts = [text for text,_ in batch_data]
+    result = batch_generate(model, tokenizer, prompts=prompts, verbose=True, format_prompts=True, max_tokens=max_tokens, temp=0.7)
+    return result
+
+async def llm_batch_processor():
+    global batch_queue
+    while True:
+        await batch_event.wait()
+        print("Batch processor triggered.")
+
+        if batch_queue:
+            batch = batch_queue[:batch_size]
+            batch_queue = batch_queue[len(batch):]
+            result = await asyncio.to_thread(batch_completion, batch)
+
+            for (_, future), res in zip(batch, result):
+                future.set_result(res)
+
+        if not batch_queue:
+            print("Batch queue is empty. Resetting event.")
+            batch_event.clear()
+
+async def async_completion(text):
+    global batch_queue
+
+    loop = asyncio.get_event_loop()
+    future = loop.create_future()
+    batch_queue.append((text, future))
+    print("Task added to batch queue.")
+
+    if len(batch_queue) >= batch_size:
+        print("Batch size reached. Triggering batch processing.")
+        batch_event.set()
+    else:
+        print(f"Task waiting for batch timeout or size threshold.")
+        await asyncio.sleep(batch_wait_time)
+        batch_event.set()
+
+    return await future

sigmaflow/clients/llm_ollama.py

@@ -0,0 +1,53 @@
+from ollama import Client, AsyncClient
+from . import *
+
+ollama_model = os.getenv('OLLAMA_MODEL')
+ollama_num_ctx = int(os.getenv('OLLAMA_NUM_CTX', 2048))
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    client = Client(timeout=6000)
+    resp = client.chat(model=ollama_model, messages=[
+        {
+            'role': 'user',
+            'content': text,
+        },
+    ], options={
+        "num_ctx": ollama_num_ctx
+    })
+
+    return resp.message.content
+
+async def async_completion(text):
+    client = AsyncClient(timeout=6000)
+
+    if (t := type(text)) is str:
+        msg = [
+            {
+                "role": "user",
+                "content": text,
+            }
+        ]
+    elif t is list:
+        msg = []
+        for role, c in zip(['user', 'assistant']*len(text), text):
+            msg.append({
+                "role": role,
+                "content": c,
+            })
+
+    async with llm_sem:
+        resp = await client.chat(
+            model=ollama_model,
+            messages=msg,
+            options={
+                "num_ctx": ollama_num_ctx
+            }
+        )
+
+    return resp.message.content

sigmaflow/clients/llm_openai.py

@@ -0,0 +1,59 @@
+from openai import OpenAI
+from openai import AsyncOpenAI
+from . import *
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    client = OpenAI(
+        api_key=os.getenv("OPENAI_API_KEY"),
+        base_url=os.getenv("OPENAI_API_URL"),
+        max_retries=20,
+        # timeout=6000.0,
+    )
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": text,
+            }
+        ],
+        model=os.getenv("OPENAI_API_MODEL")
+    )
+
+    return chat_completion.choices[0].message.content
+
+async def async_completion(text):
+    client = AsyncOpenAI(
+        api_key=os.getenv("OPENAI_API_KEY"),
+        base_url=os.getenv("OPENAI_API_URL"),
+        max_retries=20,
+        # timeout=6000.0,
+    )
+    if (t := type(text)) is str:
+        msg = [
+            {
+                "role": "user",
+                "content": text,
+            }
+        ]
+    elif t is list:
+        msg = []
+        for role, c in zip(['user', 'assistant']*len(text), text):
+            msg.append({
+                "role": role,
+                "content": c,
+            })
+
+    async with llm_sem:
+        chat_completion = await client.chat.completions.create(
+            messages=msg,
+            model=os.getenv("OPENAI_API_MODEL")
+        )
+
+    return chat_completion.choices[0].message.content

sigmaflow/clients/llm_pulse.py

@@ -0,0 +1,115 @@
+from . import *
+import json
+import httpx
+from fastapi import HTTPException
+
+pulse_model = os.getenv('PULSE_MODEL')
+pulse_token = os.getenv('PULSE_TOKEN')
+pulse_url = os.getenv('PULSE_URL')
+max_tokens = int(os.getenv('PULSE_MAX_TOKENS', 2048))
+retry_count = int(os.getenv('PULSE_RETRY_COUNT', 5))
+repetition_penalty = float(os.getenv('PULSE_REPETITION_PENALTY', 1))
+temperature = float(os.getenv('PULSE_TEMPERATURE', 0.7))
+top_p = float(os.getenv('PULSE_TOP_P', 0.1))
+top_k = int(os.getenv('PULSE_TOP_K', 10))
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    payload = json.dumps({
+        "action": "To user",
+        "parent_messages": [
+            {
+                "action": "From user",
+                "content": text
+            }
+        ],
+        "gen_kwargs": {
+            "model": pulse_model,
+            "num_return_sequences": 1,
+        }
+    })
+    headers = {
+        'accept': 'application/json',
+        'Authorization': f'Bearer {pulse_token}',
+        'Content-Type': 'application/json'
+    }
+
+    response = requests.request("POST", pulse_url, headers=headers, data=payload, verify=False)
+    if response.status_code != 200:
+        raise Exception(f"Failed to generate completion: {response.text}")
+
+    return response.json()['messages'][0]['content']['parts'][0]
+
+async def async_completion(text):
+    messages = [
+        {
+            'action': 'From user', # 'To user'
+            'content': text,
+        }
+    ]
+
+    async with llm_sem:
+        input_data = {
+            "action": "To user",
+            "parent_messages": messages,
+            "gen_kwargs": {
+                "model": pulse_model,
+                "num_return_sequences": 1,
+                "temperature": temperature,
+                "top_p": top_p,
+                "top_k": top_k,
+                "max_tokens": max_tokens,
+                "repetition_penalty": repetition_penalty,
+            }
+        }
+
+        return await _get_completion(
+            input_data=input_data,
+            retry=retry_count
+        )
+
+async def _get_completion(input_data: dict, retry: int = 5):
+    try:
+        async with httpx.AsyncClient(verify=False) as client:
+            async with client.stream(
+                "POST",
+                pulse_url,
+                json=input_data,
+                headers={
+                    "Authorization": f"Bearer {pulse_token}",
+                },
+                timeout=600,
+            ) as stream:
+                # error
+                if stream.status_code != 200:
+                    error_detail = ""
+                    try:
+                        error_detail = (await stream.aread()).decode("utf8")
+                        error_detail = str(json.loads(error_detail))
+                    except:
+                        pass
+
+                    raise HTTPException(status_code=500, detail=error_detail)
+
+                content = (await stream.aread()).decode("utf8")
+                content = json.loads(content)
+                content = "".join(content['messages'][0]['content']['parts'])
+                return content
+    except Exception as e:
+        # logger.error(json.dumps({
+        #     "retry": retry,
+        #     "input": input_data,
+        #     "error": traceback.format_exc()
+        # }, ensure_ascii=False, indent=4))
+        # 重试次数到达极限
+        if retry == 0: raise e
+
+        return await _get_completion(
+            input_data=input_data,
+            retry=retry-1,
+        )

sigmaflow/clients/llm_torch.py

@@ -0,0 +1,50 @@
+from . import *
+import torch
+from transformers import pipeline
+
+model_path = os.getenv('MODEL_PATH')
+max_new_tokens = int(os.getenv('MAX_TOKENS', 256))
+
+pipe = pipeline(
+    "text-generation", 
+    model=model_path, 
+    torch_dtype=torch.bfloat16, 
+    device_map="auto"
+)
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    if (t := type(text)) is str:
+        msg = [
+            {
+                "role": "user",
+                "content": text,
+            }
+        ]
+    elif t is list:
+        msg = []
+        for role, c in zip(['user', 'assistant']*len(text), text):
+            msg.append({
+                "role": role,
+                "content": c,
+            })
+
+    resp = pipe(
+        msg,
+        do_sample=True,
+        max_new_tokens=max_new_tokens,
+        pad_token_id = pipe.tokenizer.eos_token_id
+    )
+
+    return resp[0]["generated_text"][-1]['content']
+
+async def async_completion(text):
+    async with llm_sem:
+        resp = completion(text)
+
+    return resp

sigmaflow/clients/llm_vllm.py

@@ -0,0 +1,72 @@
+from . import *
+spec = importlib.util.find_spec('vllm')
+
+if spec:
+    from vllm import LLM, SamplingParams
+
+    model_path = os.getenv('VLLM_MODEL')
+    max_model_len = int(os.getenv('VLLM_MAX_MODEL_LEN', 256))
+    max_tokens = int(os.getenv('VLLM_MAX_TOKENS', 256))
+    batch_size = int(os.getenv('VLLM_BATCH_SIZE', 128))
+    pp = int(os.getenv('VLLM_PP', 1))
+    tp = int(os.getenv('VLLM_TP', 1))
+
+    llm = LLM(model=model_path,
+            gpu_memory_utilization=.95,
+            max_model_len=max_model_len,
+            pipeline_parallel_size=pp,
+            tensor_parallel_size=tp)
+    sampling_params = SamplingParams(temperature=0.7,
+                                    max_tokens=max_tokens,
+                                    min_p=0.15,
+                                    top_p=0.85)
+
+batch_wait_time = 1
+batch_queue = []
+batch_event = asyncio.Event()
+
+def llm_client(is_async=False):
+    if is_async:
+        return async_completion
+    else:
+        return completion
+
+def completion(text):
+    raise NotImplementedError
+
+def batch_completion(batch_data):
+    prompts = [text for text,_ in batch_data]
+    outputs = llm.generate(prompts, sampling_params)
+    result = [o.outputs[0].text for o in outputs]
+    return result
+
+async def llm_batch_processor():
+    global batch_queue
+    while True:
+        await batch_event.wait()
+
+        if batch_queue:
+            batch = batch_queue[:batch_size]
+            batch_queue = batch_queue[len(batch):]
+            result = await asyncio.to_thread(batch_completion, batch)
+
+            for (_, future), res in zip(batch, result):
+                future.set_result(res)
+
+        if not batch_queue:
+            batch_event.clear()
+
+async def async_completion(text):
+    global batch_queue
+
+    loop = asyncio.get_event_loop()
+    future = loop.create_future()
+    batch_queue.append((text, future))
+
+    if len(batch_queue) >= batch_size:
+        batch_event.set()
+    else:
+        await asyncio.sleep(batch_wait_time)
+        batch_event.set()
+
+    return await future