local-coze 0.0.1__py3-none-any.whl

local_coze/cli/image.py

@@ -0,0 +1,162 @@
+import base64
+import os
+import time
+from typing import Optional
+
+import click
+from coze_coding_utils.runtime_ctx.context import new_context
+from rich.console import Console
+
+from ..core.config import Config
+from ..image import ImageGenerationClient
+from .constants import RUN_MODE_HEADER, RUN_MODE_TEST
+
+console = Console()
+
+
+def validate_and_normalize_size(size: str) -> str:
+    if size in ["2K", "4K"]:
+        return size
+
+    try:
+        width, height = size.split("x")
+        w, h = int(width), int(height)
+
+        if 2560 <= w <= 4096 and 1440 <= h <= 4096:
+            return size
+        else:
+            console.print(
+                f"[yellow]Warning: Size {size} is out of range [2560x1440, 4096x4096], using default 2K[/yellow]"
+            )
+            return "2K"
+    except (ValueError, AttributeError):
+        console.print(
+            f"[yellow]Warning: Invalid size format {size}, using default 2K[/yellow]"
+        )
+        return "2K"
+
+
+@click.command()
+@click.option("--prompt", "-p", required=True, help="Text description of the image")
+@click.option(
+    "--output", "-o", type=click.Path(), help="Output file path"
+)
+@click.option(
+    "--size",
+    "-s",
+    default="2K",
+    help="Image size (2K, 4K, or WIDTHxHEIGHT in range [2560x1440, 4096x4096])",
+)
+@click.option(
+    "--image", "-i", help="Reference image URL or path (for image-to-image generation)"
+)
+@click.option(
+    "--images",
+    multiple=True,
+    help="Multiple reference image URLs or paths (can be used multiple times)",
+)
+@click.option("--mock", is_flag=True, help="使用 mock 模式（测试运行）")
+@click.option(
+    "--header",
+    "-H",
+    multiple=True,
+    help="自定义 HTTP 请求头 (格式: 'Key: Value' 或 'Key=Value'，可多次使用)",
+)
+@click.option("--verbose", "-v", is_flag=True, help="显示详细的 HTTP 请求日志")
+def image(
+    prompt: str,
+    output: str,
+    size: str,
+    image: Optional[str],
+    images: tuple,
+    mock: bool,
+    header: tuple,
+    verbose: bool,
+):
+    """Generate image using AI."""
+    try:
+        from .utils import parse_headers
+
+        config = Config()
+
+        ctx = None
+        custom_headers = parse_headers(header) or {}
+
+        if mock:
+            ctx = new_context(method="image.generate", headers=custom_headers)
+            custom_headers[RUN_MODE_HEADER] = RUN_MODE_TEST
+            console.print("[yellow]🧪 Mock 模式已启用（测试运行）[/yellow]")
+
+        client = ImageGenerationClient(
+            config, ctx=ctx, custom_headers=custom_headers, verbose=verbose
+        )
+
+        reference_images = None
+        if image:
+            reference_images = image
+        elif images:
+            reference_images = list(images)
+
+        validated_size = validate_and_normalize_size(size)
+
+        console.print(f"[bold cyan]Generating image...[/bold cyan]")
+        console.print(f"Prompt: [yellow]{prompt}[/yellow]")
+        console.print(f"Size: [yellow]{validated_size}[/yellow]")
+        if reference_images:
+            if isinstance(reference_images, str):
+                console.print(f"Reference image: [blue]{reference_images}[/blue]")
+            else:
+                console.print(
+                    f"Reference images: [blue]{len(reference_images)} images[/blue]"
+                )
+
+        response = client.generate(
+            prompt=prompt,
+            size=validated_size,
+            image=reference_images,
+            response_format="b64_json",
+        )
+
+        if not response.data or len(response.data) == 0:
+            raise ValueError("No image data returned")
+
+        image_data = response.data[0]
+
+        if image_data.error:
+            error_msg = image_data.error.get("message", "Unknown error")
+            raise Exception(f"Image generation failed: {error_msg}")
+
+        if image_data.b64_json:
+            if output:
+                image_bytes = base64.b64decode(image_data.b64_json)
+
+                os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
+
+                with open(output, "wb") as f:
+                    f.write(image_bytes)
+
+                console.print(f"[green]✓[/green] Image saved to: [bold]{output}[/bold]")
+            else:
+                console.print(f"[yellow]Note: Image generated but not saved (no output path specified)[/yellow]")
+        elif image_data.url:
+            console.print(f"\n[cyan]Complete Image URL:[/cyan]")
+            console.print(f"[blue]{image_data.url}[/blue]")
+            
+            if output:
+                import requests
+
+                img_response = requests.get(image_data.url)
+                img_response.raise_for_status()
+
+                os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
+
+                with open(output, "wb") as f:
+                    f.write(img_response.content)
+
+                console.print(f"[green]✓[/green] Image saved to: [bold]{output}[/bold]")
+        else:
+            raise ValueError("No image data (b64_json or url) in response")
+
+    except Exception as e:
+        console.print(f"[red]✗ Error: {str(e)}[/red]")
+        raise click.Abort()

local_coze/cli/knowledge.py

@@ -0,0 +1,195 @@
+import json
+from typing import Optional
+
+import click
+from coze_coding_utils.runtime_ctx.context import new_context
+from rich.console import Console
+from rich.table import Table
+
+from ..core.config import Config
+from ..knowledge import (
+    ChunkConfig,
+    DataSourceType,
+    KnowledgeClient,
+    KnowledgeDocument,
+)
+from .constants import RUN_MODE_HEADER, RUN_MODE_TEST
+
+console = Console()
+
+
+@click.group()
+def knowledge():
+    """Knowledge Base tools."""
+    pass
+
+
+@knowledge.command()
+@click.option("--query", "-q", required=True, help="Search query")
+@click.option("--dataset", "-d", multiple=True, help="Dataset names (tables) to search in")
+@click.option("--top-k", "-k", default=5, help="Number of results to return")
+@click.option("--min-score", "-m", default=0.0, help="Minimum similarity score")
+@click.option("--mock", is_flag=True, help="Use mock mode")
+@click.option(
+    "--header",
+    "-H",
+    multiple=True,
+    help="Custom HTTP headers (format: 'Key: Value')",
+)
+@click.option("--verbose", "-v", is_flag=True, help="Show verbose logs")
+def search(
+    query: str,
+    dataset: tuple,
+    top_k: int,
+    min_score: float,
+    mock: bool,
+    header: tuple,
+    verbose: bool,
+):
+    """Search for knowledge chunks."""
+    try:
+        from .utils import parse_headers
+
+        config = Config()
+        ctx = None
+        custom_headers = parse_headers(header) or {}
+
+        if mock:
+            ctx = new_context(method="knowledge.search", headers=custom_headers)
+            custom_headers[RUN_MODE_HEADER] = RUN_MODE_TEST
+            console.print("[yellow]🧪 Mock mode enabled[/yellow]")
+
+        client = KnowledgeClient(
+            config, ctx=ctx, custom_headers=custom_headers, verbose=verbose
+        )
+
+        dataset_list = list(dataset) if dataset else None
+
+        response = client.search(
+            query=query,
+            table_names=dataset_list,
+            top_k=top_k,
+            min_score=min_score,
+        )
+
+        if response.code != 0:
+            console.print(f"[red]Error {response.code}: {response.msg}[/red]")
+            return
+
+        table = Table(title="Search Results")
+        table.add_column("Score", style="cyan", no_wrap=True)
+        table.add_column("Content", style="green")
+        table.add_column("Chunk ID", style="dim")
+        table.add_column("Doc ID", style="dim")
+
+        for chunk in response.chunks:
+            table.add_row(
+                f"{chunk.score:.4f}",
+                chunk.content,
+                chunk.chunk_id or "",
+                chunk.doc_id or "",
+            )
+
+        console.print(table)
+
+    except Exception as e:
+        console.print(f"[red]✗ Error: {str(e)}[/red]")
+        raise click.Abort()
+
+
+@knowledge.command()
+@click.option("--dataset", "-d", required=True, help="Dataset name (table) to add to")
+@click.option("--content", "-c", multiple=True, help="Raw text content to add")
+@click.option("--url", "-u", multiple=True, help="Web URL to add")
+@click.option("--chunk-separator", default="\n", help="Chunk separator")
+@click.option("--max-tokens", default=500, help="Max tokens per chunk")
+@click.option("--remove-extra-spaces", is_flag=True, help="Normalize extra spaces")
+@click.option("--remove-urls-emails", is_flag=True, help="Strip URLs and emails")
+@click.option("--mock", is_flag=True, help="Use mock mode")
+@click.option(
+    "--header",
+    "-H",
+    multiple=True,
+    help="Custom HTTP headers (format: 'Key: Value')",
+)
+@click.option("--verbose", "-v", is_flag=True, help="Show verbose logs")
+def add(
+    dataset: str,
+    content: tuple,
+    url: tuple,
+    chunk_separator: str,
+    max_tokens: int,
+    remove_extra_spaces: bool,
+    remove_urls_emails: bool,
+    mock: bool,
+    header: tuple,
+    verbose: bool,
+):
+    """Add documents to knowledge base."""
+    try:
+        from .utils import parse_headers
+
+        config = Config()
+        ctx = None
+        custom_headers = parse_headers(header) or {}
+
+        if mock:
+            ctx = new_context(method="knowledge.add", headers=custom_headers)
+            custom_headers[RUN_MODE_HEADER] = RUN_MODE_TEST
+            console.print("[yellow]🧪 Mock mode enabled[/yellow]")
+
+        client = KnowledgeClient(
+            config, ctx=ctx, custom_headers=custom_headers, verbose=verbose
+        )
+
+        documents = []
+        for text in content:
+            documents.append(
+                KnowledgeDocument(
+                    source=DataSourceType.TEXT,
+                    raw_data=text,
+                )
+            )
+        for u in url:
+            documents.append(
+                KnowledgeDocument(
+                    source=DataSourceType.URL,
+                    url=u,
+                )
+            )
+
+        if not documents:
+            console.print("[red]No content or URL provided to add.[/red]")
+            return
+
+        chunk_config = ChunkConfig(
+            separator=chunk_separator,
+            max_tokens=max_tokens,
+            remove_extra_spaces=remove_extra_spaces,
+            remove_urls_emails=remove_urls_emails,
+        )
+
+        response = client.add_documents(
+            documents=documents,
+            table_name=dataset,
+            chunk_config=chunk_config,
+        )
+
+        if response.code != 0:
+            console.print(f"[red]Error {response.code}: {response.msg}[/red]")
+            return
+
+        table = Table(title="Add Documents Result")
+        table.add_column("Field", style="cyan", no_wrap=True)
+        table.add_column("Value", style="green")
+
+        table.add_row("Code", str(response.code))
+        table.add_row("Message", response.msg)
+        if response.doc_ids:
+            table.add_row("Doc IDs", ", ".join(response.doc_ids))
+
+        console.print(table)
+
+    except Exception as e:
+        console.print(f"[red]✗ Error: {str(e)}[/red]")
+        raise click.Abort()

local_coze/cli/search.py

@@ -0,0 +1,198 @@
+import json
+from typing import Optional
+import click
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+from rich.markdown import Markdown
+
+from ..search import SearchClient
+from ..core.config import Config
+
+console = Console()
+
+
+def display_web_results_table(response, show_url: bool = True):
+    """使用 Rich 表格显示 Web 搜索结果"""
+    if not response.web_items:
+        console.print("[yellow]没有找到搜索结果[/yellow]")
+        return
+    
+    table = Table(title="Web 搜索结果", show_header=True, header_style="bold cyan")
+    table.add_column("序号", style="dim", width=6)
+    table.add_column("标题", style="bold")
+    table.add_column("站点", style="cyan", width=15)
+    table.add_column("摘要", style="white")
+    if show_url:
+        table.add_column("URL", style="blue", width=40)
+    
+    for idx, item in enumerate(response.web_items, 1):
+        title = item.title[:50] + "..." if len(item.title) > 50 else item.title
+        site = item.site_name or "未知"
+        snippet = item.snippet[:80] + "..." if len(item.snippet) > 80 else item.snippet
+        
+        row = [str(idx), title, site, snippet]
+        if show_url:
+            url = item.url[:40] + "..." if item.url and len(item.url) > 40 else (item.url or "")
+            row.append(url)
+        
+        table.add_row(*row)
+    
+    console.print(table)
+    
+    if response.summary:
+        console.print()
+        console.print(Panel(
+            Markdown(response.summary),
+            title="[bold green]AI 摘要[/bold green]",
+            border_style="green"
+        ))
+    
+    console.print(f"\n[green]✓[/green] 找到 {len(response.web_items)} 条结果")
+
+
+def display_image_results_table(response):
+    """使用 Rich 表格显示图片搜索结果"""
+    if not response.image_items:
+        console.print("[yellow]没有找到图片结果[/yellow]")
+        return
+    
+    table = Table(title="图片搜索结果", show_header=True, header_style="bold magenta")
+    table.add_column("序号", style="dim", width=6)
+    table.add_column("标题", style="bold")
+    table.add_column("站点", style="cyan", width=15)
+    table.add_column("图片 URL", style="blue")
+    table.add_column("尺寸", style="yellow", width=12)
+    
+    for idx, item in enumerate(response.image_items, 1):
+        title = item.title[:40] + "..." if item.title and len(item.title) > 40 else (item.title or "无标题")
+        site = item.site_name or "未知"
+        image_url = item.image.url[:50] + "..." if len(item.image.url) > 50 else item.image.url
+        size = f"{item.image.width}x{item.image.height}" if item.image.width and item.image.height else "未知"
+        
+        table.add_row(str(idx), title, site, image_url, size)
+    
+    console.print(table)
+    console.print(f"\n[green]✓[/green] 找到 {len(response.image_items)} 张图片")
+
+
+def display_simple_format(response):
+    """简单文本格式输出"""
+    if response.web_items:
+        for idx, item in enumerate(response.web_items, 1):
+            console.print(f"[{idx}] {item.title} - {item.site_name or '未知'}")
+            if item.url:
+                console.print(f"    {item.url}")
+            console.print(f"    {item.snippet}")
+            console.print()
+        
+        if response.summary:
+            console.print("=" * 60)
+            console.print("AI 摘要:")
+            console.print(response.summary)
+            console.print("=" * 60)
+    
+    if response.image_items:
+        for idx, item in enumerate(response.image_items, 1):
+            console.print(f"[{idx}] {item.title or '无标题'} - {item.site_name or '未知'}")
+            console.print(f"    图片: {item.image.url}")
+            if item.image.width and item.image.height:
+                console.print(f"    尺寸: {item.image.width}x{item.image.height}")
+            console.print()
+
+
+def save_to_json(response, output_path: str):
+    """保存结果到 JSON 文件"""
+    data = response.model_dump()
+    with open(output_path, 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+    console.print(f"[green]✓[/green] 结果已保存到: {output_path}")
+
+
+@click.command()
+@click.argument("query")
+@click.option("--type", "-t", type=click.Choice(["web", "image", "web_summary"]), default="web", help="搜索类型: web(网页), image(图片), web_summary(网页+AI摘要)")
+@click.option("--count", "-c", default=10, help="返回结果数量")
+@click.option("--summary", "-s", is_flag=True, help="是否需要 AI 摘要 (仅 web 类型)")
+@click.option("--need-content", is_flag=True, help="仅返回有正文的结果 (仅 web 类型)")
+@click.option("--need-url", is_flag=True, help="仅返回有原文链接的结果 (仅 web 类型)")
+@click.option("--sites", help="指定搜索的站点范围,逗号分隔 (仅 web 类型)")
+@click.option("--block-hosts", help="屏蔽的站点,逗号分隔 (仅 web 类型)")
+@click.option("--time-range", help="发文时间范围,如: 1d, 1w, 1m (仅 web 类型)")
+@click.option("--output", "-o", type=click.Path(), help="输出 JSON 文件路径")
+@click.option("--format", "-f", type=click.Choice(["table", "json", "simple"]), default="table", help="输出格式")
+@click.option(
+    "--header",
+    "-H",
+    multiple=True,
+    help="自定义 HTTP 请求头 (格式: 'Key: Value' 或 'Key=Value'，可多次使用)",
+)
+@click.option("--verbose", "-v", is_flag=True, help="显示详细的 HTTP 请求日志")
+def search(query, type, count, summary, need_content, need_url, sites, block_hosts, time_range, output, format, header, verbose):
+    """联网搜索
+
+    支持三种搜索类型:
+    - web: 普通网页搜索 (默认)
+    - image: 图片搜索
+    - web_summary: 网页搜索 + AI 智能摘要
+
+    示例:
+        # 网页搜索
+        coze-coding-ai search "AI 最新进展"
+        coze-coding-ai search "AI 最新进展" --type web --count 20 --summary
+        
+        # 图片搜索
+        coze-coding-ai search "可爱的猫咪" --type image
+        coze-coding-ai search "可爱的猫咪" -t image -c 20
+        
+        # 网页搜索 + AI 摘要
+        coze-coding-ai search "量子计算原理" --type web_summary
+        coze-coding-ai search "量子计算原理" -t web_summary -c 15
+        
+        # 高级过滤
+        coze-coding-ai search "Python 教程" --sites "python.org,github.com" --need-content
+        coze-coding-ai search "新闻" --time-range "1d" --need-url
+    """
+    try:
+        from .utils import parse_headers
+
+        config = Config()
+        custom_headers = parse_headers(header)
+        client = SearchClient(config, custom_headers=custom_headers, verbose=verbose)
+        
+        if type == "image":
+            console.print(f"[bold magenta]正在搜索图片:[/bold magenta] {query}")
+            response = client.image_search(query=query, count=count)
+        elif type == "web_summary":
+            console.print(f"[bold cyan]正在搜索并生成摘要:[/bold cyan] {query}")
+            response = client.web_search_with_summary(query=query, count=count)
+        else:
+            console.print(f"[bold cyan]正在搜索:[/bold cyan] {query}")
+            response = client.search(
+                query=query,
+                search_type="web",
+                count=count,
+                need_content=need_content,
+                need_url=need_url,
+                sites=sites,
+                block_hosts=block_hosts,
+                need_summary=summary,
+                time_range=time_range
+            )
+        
+        if format == "json":
+            console.print_json(data=response.model_dump())
+        elif format == "simple":
+            display_simple_format(response)
+        else:
+            if type == "image":
+                display_image_results_table(response)
+            else:
+                display_web_results_table(response, show_url=True)
+        
+        if output:
+            save_to_json(response, output)
+    
+    except Exception as e:
+        console.print(f"[red]✗ 错误: {str(e)}[/red]")
+        raise click.Abort()

local_coze/cli/utils.py

@@ -0,0 +1,41 @@
+from typing import Dict, Optional, Tuple
+
+
+def parse_headers(header_strings: Tuple[str, ...]) -> Optional[Dict[str, str]]:
+    """
+    解析命令行传入的 header 字符串
+    
+    支持格式:
+    - "Key: Value"
+    - "Key=Value"
+    
+    示例:
+        --header "X-Custom: value1" --header "X-Test=value2"
+    
+    Args:
+        header_strings: 命令行传入的 header 字符串元组
+    
+    Returns:
+        Dict[str, str] or None: 解析后的 headers 字典，如果没有则返回 None
+    
+    Raises:
+        ValueError: 如果 header 格式不正确
+    """
+    if not header_strings:
+        return None
+    
+    headers = {}
+    for header_str in header_strings:
+        if ": " in header_str:
+            key, value = header_str.split(": ", 1)
+        elif "=" in header_str:
+            key, value = header_str.split("=", 1)
+        else:
+            raise ValueError(
+                f"Invalid header format: '{header_str}'. "
+                f"Use 'Key: Value' or 'Key=Value'"
+            )
+        
+        headers[key.strip()] = value.strip()
+    
+    return headers if headers else None