structai 0.1.6__py3-none-any.whl

structai/mp.py

@@ -0,0 +1,96 @@
+import concurrent.futures
+from tqdm import tqdm
+
+def multi_thread(inp_list, function, max_workers=40, use_tqdm=True):
+    results = [None] * len(inp_list)  # Initialize results list
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_index = {
+            executor.submit(function, **item): index  # Unpack item
+            for index, item in enumerate(inp_list)
+        }
+        
+        futures_iterator = concurrent.futures.as_completed(future_to_index)
+        if use_tqdm:
+            futures_iterator = tqdm(futures_iterator, total=len(future_to_index))
+
+        for future in futures_iterator:
+            index = future_to_index[future]
+            try:
+                result = future.result()
+                results[index] = result  # Place result in correct position
+            except Exception as e:
+                print(f"Error processing item {inp_list[index]}: {str(e)}")
+    
+    return results
+
+def multi_process(inp_list, function, max_workers=40, use_tqdm=True):
+    results = [None] * len(inp_list)  # Initialize results list
+    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+        future_to_index = {
+            executor.submit(function, **item): index  # Unpack item
+            for index, item in enumerate(inp_list)
+        }
+
+        futures_iterator = concurrent.futures.as_completed(future_to_index)
+        if use_tqdm:
+            futures_iterator = tqdm(futures_iterator, total=len(future_to_index))
+
+        for future in futures_iterator:
+            index = future_to_index[future]
+            try:
+                result = future.result()
+                results[index] = result  # Place result in correct position
+            except Exception as e:
+                print(f"Error processing item {inp_list[index]}: {str(e)}")
+    
+    return results
+
+if __name__ == '__main__':
+    # python -m structai.mp
+    print("Testing mp.py...")
+    import time
+    
+    def process(a, b):
+        time.sleep(0.1)
+        return a + b
+    
+    # Test data
+    inp = []
+    for a in range(10):
+        inp.append({'a': a, 'b': 100})
+    
+    # Test multi_thread
+    print("Testing multi_thread...")
+    results = multi_thread(inp, process, max_workers=5)
+    expected = [a + 100 for a in range(10)]
+    assert results == expected, f"[===ERROR===][structai][mp.py][main] multi_thread failed: {results} != {expected}"
+    print("multi_thread passed")
+    
+    # Test multi_process
+    print("Testing multi_process...")
+    results = multi_process(inp, process, max_workers=5)
+    assert results == expected, f"[===ERROR===][structai][mp.py][main] multi_process failed: {results} != {expected}"
+    print("multi_process passed")
+
+    # Test with torch if available
+    try:
+        import torch
+        if torch.cuda.is_available():
+            print("Testing with CUDA tensors...")
+            inp_cuda = []
+            for a in range(5):
+                inp_cuda.append({'a': torch.tensor(a).cuda(), 'b': torch.tensor(100).cuda()})
+            
+            # Note: multi_process with CUDA tensors might be tricky due to pickling/context
+            # We'll just test multi_thread here as it's safer for shared CUDA tensors in simple scripts
+            results = multi_thread(inp_cuda, process, max_workers=2)
+            results_cpu = [r.item() for r in results]
+            expected_cpu = [a + 100 for a in range(5)]
+            assert results_cpu == expected_cpu, f"[===ERROR===][structai][mp.py][main] CUDA multi_thread failed: {results_cpu} != {expected_cpu}"
+            print("CUDA multi_thread passed")
+        else:
+            print("Skipping CUDA tests: CUDA not available")
+    except ImportError:
+        print("Skipping torch tests: torch not installed")
+        
+    print("mp.py tests completed.")

structai/openai_server.py

@@ -0,0 +1,79 @@
+import os
+
+# =========================
+# Config
+# =========================
+
+LLM_BASE_URL = os.environ.get("LLM_BASE_URL", "")   # Do not include /v1
+LLM_API_KEY = os.environ.get("LLM_API_KEY", "")
+
+def run_server(host="0.0.0.0", port=8001):
+    """
+    Run the OpenAI proxy server.
+    """
+    from fastapi import FastAPI, Request, Response
+    import httpx
+    import uvicorn
+
+    # =========================
+    # App
+    # =========================
+
+    app = FastAPI()
+
+    # =========================
+    # Universal OpenAI Proxy
+    # =========================
+
+    @app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
+    async def openai_proxy(path: str, request: Request):
+        """
+        Universal OpenAI-compatible proxy.
+        Supports chat.completions, responses, models, embeddings, etc.
+        """
+
+        # Target URL
+        target_url = f"{LLM_BASE_URL}/v1/{path}"
+
+        # Copy headers
+        headers = dict(request.headers)
+
+        # Override authorization
+        headers["authorization"] = f"Bearer {LLM_API_KEY}"
+        headers.pop("host", None)
+
+        # Read body
+        body = await request.body()
+
+        # Forward request
+        async with httpx.AsyncClient(timeout=300) as client:
+            upstream_resp = await client.request(
+                method=request.method,
+                url=target_url,
+                headers=headers,
+                content=body,
+                params=request.query_params,
+            )
+
+        # -------------------------
+        # IMPORTANT: header cleanup
+        # -------------------------
+        resp_headers = dict(upstream_resp.headers)
+
+        # httpx already decompressed the body
+        # Do NOT forward these headers
+        resp_headers.pop("content-encoding", None)
+        resp_headers.pop("content-length", None)
+
+        return Response(
+            content=upstream_resp.content,
+            status_code=upstream_resp.status_code,
+            headers=resp_headers,
+            media_type=resp_headers.get("content-type"),
+        )
+
+    uvicorn.run(app, host=host, port=port)
+
+if __name__ == "__main__":
+    # python -m structai.openai_server
+    run_server()

structai/utils.py

@@ -0,0 +1,209 @@
+from functools import wraps
+import threading
+import os
+
+
+def run_with_timeout(func, args=(), kwargs=None, timeout=None):
+    """
+    Run a function with a timeout limit.
+    
+    Args:
+        func (callable): The function to run.
+        args (tuple): Positional arguments for the function.
+        kwargs (dict): Keyword arguments for the function.
+        timeout (float or None): Maximum allowed execution time in seconds.
+    
+    Returns:
+        The function's return value.
+    
+    Raises:
+        TimeoutError: If execution exceeds the allowed time.
+        Exception: Any exception raised by the function.
+    """
+    if kwargs is None:
+        kwargs = {}
+        
+    if timeout is None:
+        return func(*args, **kwargs)
+
+    result = [None]
+    exc = [None]
+
+    def target():
+        try:
+            result[0] = func(*args, **kwargs)
+        except Exception as e:
+            exc[0] = e
+
+    thread = threading.Thread(target=target)
+    thread.start()
+    thread.join(timeout)
+
+    if thread.is_alive():
+        raise TimeoutError(f"[{str(func.__name__)} timeout after {timeout}s]")
+
+    if exc[0] is not None:
+        raise exc[0]
+
+    return result[0]
+
+
+def timeout_limit(timeout=None):
+    """
+    A decorator for enforcing function execution time limits.
+
+    Args:
+        timeout (float or None): Maximum allowed execution time in seconds.
+                                If None, no time limit is applied.
+
+    Returns:
+        The function's return value if completed in time.
+        Raises TimeoutError if execution exceeds the allowed time.
+    """
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return run_with_timeout(func, args, kwargs, timeout)
+        return wrapper
+
+    return decorator
+
+
+def remove_tag(s: str, tags: list[str] = ["<think>", "</think>", "<answer>", "</answer>"], r: str = "\n"):
+    for tag in tags:
+        s = s.replace(tag, r)
+    return s.strip(r).strip()
+
+
+def parse_think_answer(text):
+    if not isinstance(text, str):
+        raise ValueError("Input text must be a string")
+    
+    text = text.strip()
+    
+    tag_think_start = "<think>"
+    tag_think_end = "</think>"
+    tag_answer_start = "<answer>"
+    tag_answer_end = "</answer>"
+    
+    idx_answer_start = text.find(tag_answer_start)
+    idx_think_end = text.find(tag_think_end)
+    
+    think_raw = ""
+    answer_raw = ""
+    
+    if idx_answer_start != -1:
+        think_raw = text[:idx_answer_start]
+        answer_raw = text[idx_answer_start + len(tag_answer_start):]
+        
+        idx_answer_end = answer_raw.find(tag_answer_end)
+        if idx_answer_end != -1:
+            answer_raw = answer_raw[:idx_answer_end]
+            
+    elif idx_think_end != -1:
+        think_raw = text[:idx_think_end]
+        answer_raw = text[idx_think_end + len(tag_think_end):]
+        
+        idx_answer_end = answer_raw.find(tag_answer_end)
+        if idx_answer_end != -1:
+            answer_raw = answer_raw[:idx_answer_end]
+            
+    else:
+        raise ValueError("Could not parse response: missing <answer> or </think> tags.")
+        
+    idx_think_start = think_raw.find(tag_think_start)
+    if idx_think_start != -1:
+        think_raw = think_raw[idx_think_start + len(tag_think_start):]
+
+    think = remove_tag(think_raw, [tag_think_start, tag_think_end, tag_answer_start, tag_answer_end], "")
+    answer = remove_tag(answer_raw, [tag_think_start, tag_think_end, tag_answer_start, tag_answer_end], "")
+    
+    if not think:
+        raise ValueError("Parsed think content is empty")
+    if not answer:
+        raise ValueError("Parsed answer content is empty")
+        
+    return think, answer
+
+
+def extract_within_tags(content: str, start_tag='<answer>', end_tag='</answer>', default_return=None):
+    content = str(content)
+    start_index = content.rfind(start_tag)
+    if start_index != -1:
+        end_index = content.find(end_tag, start_index)
+        if end_index != -1:
+            return content[start_index + len(start_tag):end_index].strip()
+    return default_return
+
+
+def get_all_file_paths(directory, suffix=''):
+    """
+    Get all file paths with the specified suffix in the directory.
+    """
+    file_paths = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.endswith(suffix):
+                file_paths.append(os.path.join(root, file))
+    return file_paths
+
+
+if __name__ == "__main__":
+    # python -m structai.utils
+    print("Testing utils.py...")
+    import time
+    
+    # Test run_with_timeout
+    print("Testing run_with_timeout...")
+    def slow_func(seconds):
+        time.sleep(seconds)
+        return "done"
+        
+    # Should pass
+    res = run_with_timeout(slow_func, args=(0.1,), timeout=1.0)
+    assert res == "done", f"[===ERROR===][structai][utils.py][main] run_with_timeout failed: {res} != done"
+    
+    # Should fail
+    try:
+        run_with_timeout(slow_func, args=(0.5,), timeout=0.1)
+        print("run_with_timeout failed to timeout (unexpected)")
+    except TimeoutError:
+        print("run_with_timeout correctly timed out")
+        
+    # Test timeout_limit decorator
+    print("Testing timeout_limit decorator...")
+    @timeout_limit(timeout=0.2)
+    def decorated_slow_func(seconds):
+        time.sleep(seconds)
+        return "done"
+        
+    try:
+        decorated_slow_func(0.5)
+        print("timeout_limit failed to timeout (unexpected)")
+    except TimeoutError:
+        print("timeout_limit correctly timed out")
+        
+    # Test remove_tag
+    print("Testing remove_tag...")
+    text = "<think>thinking...</think> <answer>42</answer>"
+    cleaned = remove_tag(text)
+    assert "thinking..." in cleaned and "42" in cleaned, f"[===ERROR===][structai][utils.py][main] remove_tag failed: {cleaned}"
+    assert "<think>" not in cleaned, f"[===ERROR===][structai][utils.py][main] remove_tag failed: {cleaned}"
+    print("remove_tag passed")
+    
+    # Test parse_think_answer
+    print("Testing parse_think_answer...")
+    text = "<think>This is thinking</think><answer>This is answer</answer>"
+    think, answer = parse_think_answer(text)
+    assert think == "This is thinking", f"[===ERROR===][structai][utils.py][main] parse_think_answer failed: {think} != This is thinking"
+    assert answer == "This is answer", f"[===ERROR===][structai][utils.py][main] parse_think_answer failed: {answer} != This is answer"
+    print("parse_think_answer passed")
+    
+    # Test extract_within_tags
+    print("Testing extract_within_tags...")
+    text = "prefix <tag>content</tag> suffix"
+    extracted = extract_within_tags(text, "<tag>", "</tag>")
+    assert extracted == "content", f"[===ERROR===][structai][utils.py][main] extract_within_tags failed: {extracted} != content"
+    print("extract_within_tags passed")
+    
+    print("utils.py tests completed.")