mirror of
https://github.com/GNS3/gns3-server.git
synced 2026-08-27 20:40:13 +03:00
feat(agent): remove checkpoint utilities from public API
Removed checkpoint-related imports and exports from the agent package's __init__.py to clean up the public API. This change reduces unnecessary exposure of internal checkpoint utility functions, focusing the public interface on the core agent_builder functionality.
This commit is contained in:
YueGuobin
parent
cabcf4f1cb
commit
15dc1e7ea0
@ -4,13 +4,6 @@ FlowNet-Lab Agent Package
|
||||
This package contains the main FlowNet-Lab agent implementation for network automation tasks.
|
||||
"""
|
||||
|
||||
from .checkpoint_utils import (
|
||||
export_checkpoint_to_file,
|
||||
generate_thread_id,
|
||||
import_checkpoint_from_file,
|
||||
list_thread_ids,
|
||||
validate_checkpoint_data,
|
||||
)
|
||||
from .gns3_copilot import agent_builder
|
||||
|
||||
# Dynamic version management
|
||||
@ -27,9 +20,4 @@ __url__ = "https://github.com/yueguobin/gns3-copilot"
|
||||
|
||||
__all__ = [
|
||||
"agent_builder",
|
||||
"list_thread_ids",
|
||||
"generate_thread_id",
|
||||
"validate_checkpoint_data",
|
||||
"export_checkpoint_to_file",
|
||||
"import_checkpoint_from_file",
|
||||
]
|
||||
|
||||
@ -1,625 +0,0 @@
|
||||
"""
|
||||
FlowNet-Lab Checkpoint Utilities
|
||||
|
||||
This module provides utility functions for interacting with LangGraph checkpoint
|
||||
database, including thread ID listing, checkpoint export, import, validation,
|
||||
and session inspection.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from langchain.messages import AIMessage, HumanMessage, ToolMessage
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.types import RunnableConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def list_thread_ids(checkpointer: Any) -> list[str]:
|
||||
"""
|
||||
Get all unique thread IDs from LangGraph checkpoint database.
|
||||
|
||||
Args:
|
||||
checkpointer: LangGraph checkpointer instance.
|
||||
|
||||
Returns:
|
||||
list: List of unique thread IDs ordered by most recent activity.
|
||||
Returns empty list on error or if table doesn't exist.
|
||||
"""
|
||||
try:
|
||||
res = checkpointer.conn.execute(
|
||||
"SELECT DISTINCT thread_id FROM checkpoints ORDER BY rowid DESC"
|
||||
).fetchall()
|
||||
return [r[0] for r in res]
|
||||
except Exception as e:
|
||||
# Table might not exist yet, return empty list
|
||||
logger.debug("Error listing thread IDs (table may not exist): %s", e)
|
||||
return []
|
||||
|
||||
|
||||
def generate_thread_id() -> str:
|
||||
"""
|
||||
Generate a new unique thread ID.
|
||||
|
||||
Returns:
|
||||
str: A UUID-based thread ID.
|
||||
"""
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def validate_checkpoint_data(data: Any) -> tuple[bool, str]:
|
||||
"""
|
||||
Validate checkpoint data structure.
|
||||
|
||||
Ensures the imported checkpoint data has the required structure for
|
||||
importing into LangGraph checkpointer.
|
||||
|
||||
Args:
|
||||
data: Dictionary containing checkpoint data.
|
||||
|
||||
Returns:
|
||||
tuple: (is_valid, error_message)
|
||||
- is_valid: True if data is valid, False otherwise
|
||||
- error_message: Empty string if valid, error description if invalid
|
||||
"""
|
||||
# Check if data is a dictionary
|
||||
if not isinstance(data, dict):
|
||||
return False, "Data must be a dictionary"
|
||||
|
||||
# Check if checkpoint field exists
|
||||
if "checkpoint" not in data:
|
||||
return False, "Missing required field: checkpoint"
|
||||
|
||||
checkpoint = data["checkpoint"]
|
||||
|
||||
# Check if checkpoint is a dictionary
|
||||
if not isinstance(checkpoint, dict):
|
||||
return False, "checkpoint must be a dictionary"
|
||||
|
||||
# Check required top-level fields in checkpoint
|
||||
required_fields = ["v", "ts", "id", "channel_values", "channel_versions"]
|
||||
if not all(field in checkpoint for field in required_fields):
|
||||
missing = [f for f in required_fields if f not in checkpoint]
|
||||
return False, f"Missing required checkpoint field: {missing[0]}"
|
||||
|
||||
# Check if channel_values is a dictionary
|
||||
channel_values = checkpoint["channel_values"]
|
||||
if not isinstance(channel_values, dict):
|
||||
return False, "channel_values must be a dictionary"
|
||||
|
||||
# Check if messages field exists in channel_values
|
||||
if "messages" not in channel_values:
|
||||
return False, "Missing required field: channel_values.messages"
|
||||
|
||||
# Check if messages is a list
|
||||
if not isinstance(channel_values["messages"], list):
|
||||
return False, "channel_values.messages must be a list"
|
||||
|
||||
# All validations passed
|
||||
return True, ""
|
||||
|
||||
|
||||
def parse_message_string(msg_str: str) -> dict:
|
||||
"""
|
||||
Parse a message string representation back to a dictionary.
|
||||
|
||||
Handles the format: "content='xxx' additional_kwargs={} response_metadata={}"
|
||||
|
||||
Args:
|
||||
msg_str: String representation of a LangChain message.
|
||||
|
||||
Returns:
|
||||
dict: Parsed message data with type and content.
|
||||
"""
|
||||
import re
|
||||
|
||||
# Try to determine message type from string content
|
||||
msg_type = "unknown"
|
||||
|
||||
# Check for ToolMessage pattern (has tool_call_id or tool_name)
|
||||
if "tool_call_id=" in msg_str or "name=" in msg_str:
|
||||
msg_type = "tool"
|
||||
# Check for AIMessage pattern (has tool_calls or is an AI response)
|
||||
elif "tool_calls=" in msg_str:
|
||||
msg_type = "ai"
|
||||
else:
|
||||
# Default to human if no other indicators
|
||||
msg_type = "human"
|
||||
|
||||
# Extract content using regex
|
||||
content_match = re.search(r"content='([^']*)'|content=\"([^\"]*)\"", msg_str)
|
||||
content = (
|
||||
content_match.group(1)
|
||||
if content_match and content_match.group(1)
|
||||
else (
|
||||
content_match.group(2) if content_match and content_match.group(2) else ""
|
||||
)
|
||||
)
|
||||
|
||||
return {"type": msg_type, "content": content, "original_string": msg_str}
|
||||
|
||||
|
||||
def serialize_message(msg: Any) -> dict:
|
||||
"""
|
||||
Serialize a LangChain message to a dictionary for JSON storage.
|
||||
|
||||
Ensures all message fields are properly serialized for UI compatibility,
|
||||
including tool_calls structure for AIMessage and all metadata fields.
|
||||
|
||||
Args:
|
||||
msg: LangChain message object (HumanMessage, AIMessage, or ToolMessage)
|
||||
or a string representation of a message.
|
||||
|
||||
Returns:
|
||||
dict: Serialized message with type, content, and all metadata.
|
||||
"""
|
||||
# If it's already a string (stored in database), parse it
|
||||
if isinstance(msg, str):
|
||||
return parse_message_string(msg)
|
||||
|
||||
if isinstance(msg, HumanMessage):
|
||||
return {
|
||||
"type": "human",
|
||||
"content": msg.content,
|
||||
"additional_kwargs": msg.additional_kwargs,
|
||||
"response_metadata": msg.response_metadata,
|
||||
"id": msg.id,
|
||||
}
|
||||
elif isinstance(msg, AIMessage):
|
||||
# Serialize tool_calls with complete structure
|
||||
tool_calls = []
|
||||
if msg.tool_calls:
|
||||
for tool_call in msg.tool_calls:
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": tool_call.get("id", ""),
|
||||
"name": tool_call.get("name", ""),
|
||||
"args": tool_call.get("args", {}),
|
||||
"type": tool_call.get("type", "tool_call"),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"type": "ai",
|
||||
"content": msg.content,
|
||||
"additional_kwargs": msg.additional_kwargs,
|
||||
"response_metadata": msg.response_metadata,
|
||||
"tool_calls": tool_calls,
|
||||
"id": msg.id,
|
||||
}
|
||||
elif isinstance(msg, ToolMessage):
|
||||
return {
|
||||
"type": "tool",
|
||||
"content": msg.content,
|
||||
"tool_call_id": msg.tool_call_id,
|
||||
"name": msg.name,
|
||||
"additional_kwargs": msg.additional_kwargs,
|
||||
"response_metadata": msg.response_metadata,
|
||||
"id": msg.id,
|
||||
}
|
||||
else:
|
||||
# For any other message types, try to serialize as dict
|
||||
return {"type": "unknown", "content": str(msg)}
|
||||
|
||||
|
||||
def deserialize_message(msg_dict: dict) -> Any:
|
||||
"""
|
||||
Deserialize a dictionary back to a LangChain message object.
|
||||
|
||||
Ensures proper reconstruction of message objects with all required fields
|
||||
for UI compatibility, including tool_calls structure for AIMessage.
|
||||
|
||||
Args:
|
||||
msg_dict: Dictionary containing serialized message data.
|
||||
|
||||
Returns:
|
||||
LangChain message object (HumanMessage, AIMessage, or ToolMessage).
|
||||
|
||||
Raises:
|
||||
ValueError: If message type is unknown or required fields are missing.
|
||||
"""
|
||||
msg_type = msg_dict.get("type", "unknown")
|
||||
|
||||
if msg_type == "human":
|
||||
return HumanMessage(
|
||||
content=msg_dict.get("content", ""),
|
||||
additional_kwargs=msg_dict.get("additional_kwargs", {}),
|
||||
response_metadata=msg_dict.get("response_metadata", {}),
|
||||
id=msg_dict.get("id"),
|
||||
)
|
||||
elif msg_type == "ai":
|
||||
# Reconstruct tool_calls with proper structure
|
||||
tool_calls = []
|
||||
serialized_tool_calls = msg_dict.get("tool_calls", [])
|
||||
|
||||
if serialized_tool_calls:
|
||||
for tool_call in serialized_tool_calls:
|
||||
# Handle both dict and list formats
|
||||
if isinstance(tool_call, dict):
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": tool_call.get("id", ""),
|
||||
"name": tool_call.get("name", ""),
|
||||
"args": tool_call.get("args", {}),
|
||||
"type": tool_call.get("type", "tool_call"),
|
||||
}
|
||||
)
|
||||
|
||||
return AIMessage(
|
||||
content=msg_dict.get("content", ""),
|
||||
additional_kwargs=msg_dict.get("additional_kwargs", {}),
|
||||
response_metadata=msg_dict.get("response_metadata", {}),
|
||||
tool_calls=tool_calls,
|
||||
id=msg_dict.get("id"),
|
||||
)
|
||||
elif msg_type == "tool":
|
||||
return ToolMessage(
|
||||
content=msg_dict.get("content", ""),
|
||||
tool_call_id=msg_dict.get("tool_call_id", ""),
|
||||
name=msg_dict.get("name", ""),
|
||||
additional_kwargs=msg_dict.get("additional_kwargs", {}),
|
||||
response_metadata=msg_dict.get("response_metadata", {}),
|
||||
id=msg_dict.get("id"),
|
||||
)
|
||||
else:
|
||||
# Return as dict for unknown types
|
||||
logger.warning("Unknown message type: %s, returning dict", msg_type)
|
||||
return msg_dict
|
||||
|
||||
|
||||
def export_checkpoint_to_file(
|
||||
checkpointer: Any, thread_id: str, file_path: str
|
||||
) -> bool:
|
||||
"""
|
||||
Export checkpoint data to a .txt file in JSON format.
|
||||
|
||||
Exports the checkpoint data for a specific thread to a text file.
|
||||
The exported data includes the complete checkpoint state with messages,
|
||||
conversation title, config, and metadata. Messages are properly serialized
|
||||
to ensure they can be correctly deserialized on import.
|
||||
|
||||
Args:
|
||||
checkpointer: LangGraph checkpointer instance.
|
||||
thread_id: The thread ID to export.
|
||||
file_path: Path to the output .txt file.
|
||||
|
||||
Returns:
|
||||
bool: True if export succeeded, False otherwise.
|
||||
"""
|
||||
try:
|
||||
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
|
||||
checkpoint_tuple = checkpointer.get_tuple(config)
|
||||
|
||||
if checkpoint_tuple is None:
|
||||
logger.error("Checkpoint not found for thread_id: %s", thread_id)
|
||||
return False
|
||||
|
||||
# Create a copy of checkpoint data to avoid modifying the original
|
||||
checkpoint_data = dict(checkpoint_tuple.checkpoint)
|
||||
|
||||
# Serialize messages for proper JSON export
|
||||
if (
|
||||
"channel_values" in checkpoint_data
|
||||
and "messages" in checkpoint_data["channel_values"]
|
||||
):
|
||||
messages = checkpoint_data["channel_values"]["messages"]
|
||||
serialized_messages = [serialize_message(msg) for msg in messages]
|
||||
checkpoint_data["channel_values"]["messages"] = serialized_messages
|
||||
|
||||
# Save complete checkpoint data including config and metadata
|
||||
export_data = {
|
||||
"checkpoint": checkpoint_data,
|
||||
"config": checkpoint_tuple.config,
|
||||
"metadata": checkpoint_tuple.metadata,
|
||||
}
|
||||
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
json.dump(export_data, f, ensure_ascii=False, indent=2, default=str)
|
||||
|
||||
logger.info("Checkpoint exported to %s for thread_id: %s", file_path, thread_id)
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to export checkpoint: %s", e)
|
||||
return False
|
||||
|
||||
|
||||
def validate_messages_for_ui(messages: list) -> tuple[bool, str, list[str]]:
|
||||
"""
|
||||
Validate that messages can be correctly rendered by the UI.
|
||||
|
||||
Checks that all messages are proper LangChain message objects with
|
||||
required fields for UI rendering (chat.py compatibility).
|
||||
|
||||
Args:
|
||||
messages: List of message objects to validate.
|
||||
|
||||
Returns:
|
||||
tuple: (is_valid, error_message, validation_errors)
|
||||
- is_valid: True if all messages are valid for UI
|
||||
- error_message: Summary error message
|
||||
- validation_errors: List of specific validation errors per message
|
||||
"""
|
||||
validation_errors = []
|
||||
|
||||
if not messages:
|
||||
return True, "", []
|
||||
|
||||
for idx, msg in enumerate(messages):
|
||||
msg_error = f"Message {idx}: "
|
||||
|
||||
# Check if message is a recognized type
|
||||
if isinstance(msg, HumanMessage):
|
||||
# HumanMessage requires content
|
||||
if not hasattr(msg, "content") or msg.content is None:
|
||||
validation_errors.append(msg_error + "Missing content field")
|
||||
continue
|
||||
elif isinstance(msg, AIMessage):
|
||||
# AIMessage should have content and tool_calls
|
||||
if not hasattr(msg, "content"):
|
||||
validation_errors.append(msg_error + "Missing content field")
|
||||
|
||||
# Validate tool_calls if present
|
||||
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
||||
for tool_idx, tool_call in enumerate(msg.tool_calls):
|
||||
required_fields = ["id", "name", "args"]
|
||||
missing = [f for f in required_fields if f not in tool_call]
|
||||
if missing:
|
||||
validation_errors.append(
|
||||
f"{msg_error} Tool call {tool_idx} missing: {', '.join(missing)}"
|
||||
)
|
||||
continue
|
||||
elif isinstance(msg, ToolMessage):
|
||||
# ToolMessage requires content, tool_call_id, and name
|
||||
if not hasattr(msg, "content"):
|
||||
validation_errors.append(msg_error + "Missing content field")
|
||||
if not hasattr(msg, "tool_call_id") or not msg.tool_call_id:
|
||||
validation_errors.append(msg_error + "Missing or empty tool_call_id")
|
||||
if not hasattr(msg, "name") or not msg.name:
|
||||
validation_errors.append(msg_error + "Missing or empty name")
|
||||
continue
|
||||
else:
|
||||
validation_errors.append(
|
||||
f"{msg_error} Unknown message type: {type(msg).__name__}"
|
||||
)
|
||||
continue
|
||||
|
||||
is_valid = len(validation_errors) == 0
|
||||
error_message = "; ".join(validation_errors) if validation_errors else ""
|
||||
|
||||
return is_valid, error_message, validation_errors
|
||||
|
||||
|
||||
def inspect_session(
|
||||
thread_id: str, graph: Pregel, verbose: bool = False
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Inspect and return human-readable session state using graph.get_state().
|
||||
|
||||
Provides detailed information about a session including message statistics,
|
||||
UI compatibility, and current execution state.
|
||||
|
||||
Args:
|
||||
thread_id: Thread ID to inspect.
|
||||
graph: Compiled LangGraph agent instance.
|
||||
verbose: If True, include detailed message contents in output.
|
||||
|
||||
Returns:
|
||||
dict: Human-readable session information including:
|
||||
- next: Next action to be executed
|
||||
- message_count: Number of messages
|
||||
- message_types: Dictionary counting message types
|
||||
- latest_message: Content of latest message
|
||||
- step: Current step number
|
||||
- pending_tasks: Number of pending tasks
|
||||
- has_interrupts: Whether there are interrupts
|
||||
- conversation_title: Session title
|
||||
- selected_project: Currently selected GNS3 project
|
||||
- ui_compatible: Whether messages are compatible with UI
|
||||
- validation_errors: List of validation errors (if any)
|
||||
- messages_preview: Preview of messages (if verbose=True)
|
||||
"""
|
||||
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
|
||||
|
||||
try:
|
||||
snapshot = graph.get_state(config)
|
||||
|
||||
# Extract message information
|
||||
messages = snapshot.values.get("messages", [])
|
||||
message_count = len(messages)
|
||||
|
||||
# Count message types
|
||||
message_types = {"human": 0, "ai": 0, "tool": 0, "unknown": 0}
|
||||
for msg in messages:
|
||||
if isinstance(msg, HumanMessage):
|
||||
message_types["human"] += 1
|
||||
elif isinstance(msg, AIMessage):
|
||||
message_types["ai"] += 1
|
||||
elif isinstance(msg, ToolMessage):
|
||||
message_types["tool"] += 1
|
||||
else:
|
||||
message_types["unknown"] += 1
|
||||
|
||||
# Get latest message content
|
||||
latest_message = None
|
||||
if messages:
|
||||
latest_msg = messages[-1]
|
||||
if hasattr(latest_msg, "content"):
|
||||
if isinstance(latest_msg.content, str):
|
||||
latest_message = latest_msg.content
|
||||
elif isinstance(latest_msg.content, list) and latest_msg.content:
|
||||
# Handle Gemini format (list with text field)
|
||||
if isinstance(latest_msg.content[0], dict):
|
||||
latest_message = latest_msg.content[0].get(
|
||||
"text", str(latest_msg.content)
|
||||
)
|
||||
else:
|
||||
latest_message = str(latest_msg.content)
|
||||
else:
|
||||
latest_message = str(latest_msg.content)
|
||||
|
||||
# Validate UI compatibility
|
||||
is_valid, error_msg, validation_errors = validate_messages_for_ui(messages)
|
||||
|
||||
# Build result dictionary
|
||||
result = {
|
||||
"thread_id": thread_id,
|
||||
"next": snapshot.next,
|
||||
"message_count": message_count,
|
||||
"message_types": message_types,
|
||||
"latest_message": latest_message,
|
||||
"step": snapshot.metadata.get("step", "N/A")
|
||||
if snapshot.metadata
|
||||
else "N/A",
|
||||
"pending_tasks": len(snapshot.tasks),
|
||||
"has_interrupts": len(snapshot.interrupts) > 0,
|
||||
"conversation_title": snapshot.values.get("conversation_title"),
|
||||
"selected_project": snapshot.values.get("selected_project"),
|
||||
"ui_compatible": is_valid,
|
||||
"validation_error": error_msg,
|
||||
"validation_errors": validation_errors,
|
||||
}
|
||||
|
||||
# Add verbose details if requested
|
||||
if verbose:
|
||||
messages_preview = []
|
||||
for idx, msg in enumerate(messages):
|
||||
msg_preview = {
|
||||
"index": idx,
|
||||
"type": type(msg).__name__,
|
||||
}
|
||||
if hasattr(msg, "content"):
|
||||
msg_preview["content"] = str(msg.content)[
|
||||
:200
|
||||
] # Truncate long content
|
||||
if (
|
||||
isinstance(msg, AIMessage)
|
||||
and hasattr(msg, "tool_calls")
|
||||
and msg.tool_calls
|
||||
):
|
||||
msg_preview["tool_calls_count"] = len(msg.tool_calls)
|
||||
messages_preview.append(msg_preview)
|
||||
result["messages_preview"] = messages_preview
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to inspect session %s: %s", thread_id, e)
|
||||
return {
|
||||
"thread_id": thread_id,
|
||||
"error": str(e),
|
||||
"message_count": 0,
|
||||
"message_types": {"human": 0, "ai": 0, "tool": 0, "unknown": 0},
|
||||
"ui_compatible": False,
|
||||
"validation_error": f"Failed to get state: {e}",
|
||||
}
|
||||
|
||||
|
||||
def import_checkpoint_from_file(
|
||||
checkpointer: Any, file_path: str, new_thread_id: str | None = None
|
||||
) -> tuple[bool, str]:
|
||||
"""
|
||||
Import checkpoint data from a .txt file to a new thread.
|
||||
|
||||
Reads checkpoint data from a JSON-formatted .txt file and imports it
|
||||
into a new thread in the LangGraph checkpointer. The imported data is
|
||||
validated and messages are deserialized before insertion.
|
||||
|
||||
Args:
|
||||
checkpointer: LangGraph checkpointer instance.
|
||||
file_path: Path to .txt file containing checkpoint data.
|
||||
new_thread_id: Optional thread ID for the new thread.
|
||||
If None, a new UUID will be generated.
|
||||
|
||||
Returns:
|
||||
tuple: (success, result)
|
||||
- success: True if import succeeded, False otherwise
|
||||
- result: New thread ID if success, error message if failed
|
||||
"""
|
||||
try:
|
||||
with open(file_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Validate checkpoint data structure
|
||||
is_valid, error_msg = validate_checkpoint_data(data)
|
||||
if not is_valid:
|
||||
logger.error("Invalid checkpoint data: %s", error_msg)
|
||||
return False, error_msg
|
||||
|
||||
# Generate new thread ID if not provided
|
||||
if new_thread_id is None:
|
||||
new_thread_id = generate_thread_id()
|
||||
|
||||
# Create a copy of checkpoint data to avoid modifying original
|
||||
checkpoint_data = dict(data["checkpoint"])
|
||||
|
||||
# Deserialize messages if they are in serialized format
|
||||
if (
|
||||
"channel_values" in checkpoint_data
|
||||
and "messages" in checkpoint_data["channel_values"]
|
||||
):
|
||||
messages = checkpoint_data["channel_values"]["messages"]
|
||||
|
||||
# Check if messages are in serialized format (have 'type' field)
|
||||
if messages and isinstance(messages[0], dict) and "type" in messages[0]:
|
||||
deserialized_messages = []
|
||||
for msg_dict in messages:
|
||||
try:
|
||||
deserialized_msg = deserialize_message(msg_dict)
|
||||
deserialized_messages.append(deserialized_msg)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to deserialize message: %s. Error: %s",
|
||||
msg_dict.get("type", "unknown"),
|
||||
e,
|
||||
)
|
||||
# Skip invalid messages or add as dict
|
||||
deserialized_messages.append(msg_dict)
|
||||
checkpoint_data["channel_values"]["messages"] = deserialized_messages
|
||||
|
||||
# Rebuild complete config with all required keys
|
||||
saved_config = data.get("config", {})
|
||||
new_config = {
|
||||
**saved_config,
|
||||
"configurable": {
|
||||
**saved_config.get("configurable", {}),
|
||||
"thread_id": new_thread_id,
|
||||
"checkpoint_ns": "", # Required: checkpoint namespace (empty string)
|
||||
"checkpoint_id": str(uuid.uuid4()), # Required: new checkpoint ID
|
||||
},
|
||||
}
|
||||
|
||||
# Use saved metadata if available, otherwise create new
|
||||
metadata = data.get("metadata", {"source": "import"})
|
||||
if "source" not in metadata:
|
||||
metadata["source"] = "import"
|
||||
|
||||
new_versions = checkpoint_data["channel_versions"]
|
||||
|
||||
checkpointer.put(
|
||||
config=new_config,
|
||||
checkpoint=checkpoint_data,
|
||||
metadata=metadata,
|
||||
new_versions=new_versions,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Checkpoint imported from %s to new thread_id: %s", file_path, new_thread_id
|
||||
)
|
||||
return True, new_thread_id
|
||||
|
||||
except FileNotFoundError:
|
||||
error_msg = f"File not found: {file_path}"
|
||||
logger.error(error_msg)
|
||||
return False, error_msg
|
||||
except json.JSONDecodeError as e:
|
||||
error_msg = f"Invalid JSON format in file: {e}"
|
||||
logger.error(error_msg)
|
||||
return False, error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to import checkpoint: {e}"
|
||||
logger.error(error_msg)
|
||||
return False, error_msg
|
||||
@ -1,235 +0,0 @@
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
#
|
||||
# This file is part of FlowNet-Lab.
|
||||
#
|
||||
# FlowNet-Lab is free software: you can redistribute it and/or modify it
|
||||
# under the terms of the GNU Affero General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# FlowNet-Lab is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License along
|
||||
# with FlowNet-Lab. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""
|
||||
Experiment Deployment Agent
|
||||
|
||||
A LangGraph-based agent for automated GNS3 experiment deployment.
|
||||
Parses experiment plans and deploys complete GNS3 lab environments.
|
||||
"""
|
||||
|
||||
import operator
|
||||
from typing import Annotated, TypedDict, Any
|
||||
|
||||
from langchain.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage, ToolMessage
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.managed.is_last_step import RemainingSteps
|
||||
|
||||
import logging
|
||||
|
||||
from gns3_copilot.agent.model_factory import create_base_model_with_tools
|
||||
from gns3_copilot.prompts.experiment_deploy_prompt import EXPERIMENT_DEPLOY_SYSTEM_PROMPT
|
||||
from gns3_copilot.tools_v2 import (
|
||||
GNS3TemplateTool,
|
||||
GNS3CreateNodeTool,
|
||||
GNS3LinkTool,
|
||||
GNS3StartNodeQuickTool,
|
||||
GNS3UpdateNodeNameTool,
|
||||
)
|
||||
from gns3_copilot.gns3_client import GNS3ProjectCreate, GNS3ProjectList, GNS3TopologyTool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Create tool instances - use existing GNS3ProjectCreate directly
|
||||
tools = [
|
||||
GNS3ProjectList(), # List existing projects to avoid duplicates
|
||||
GNS3TemplateTool(),
|
||||
GNS3ProjectCreate(), # Use existing tool directly
|
||||
GNS3CreateNodeTool(),
|
||||
GNS3TopologyTool(), # Get topology to see actual port names
|
||||
GNS3LinkTool(),
|
||||
GNS3UpdateNodeNameTool(),
|
||||
GNS3StartNodeQuickTool(), # Use quick start to avoid HTTP timeouts
|
||||
]
|
||||
|
||||
tools_by_name = {tool.name: tool for tool in tools}
|
||||
|
||||
logger.info(f"Experiment Deploy Agent initialized with {len(tools)} tools")
|
||||
logger.debug(f"Available tools: {[tool.name for tool in tools]}")
|
||||
|
||||
|
||||
# Define state
|
||||
class ExperimentDeployState(TypedDict):
|
||||
"""
|
||||
Experiment deployment agent state.
|
||||
|
||||
Attributes:
|
||||
messages: Conversation messages
|
||||
llm_calls: Number of LLM calls made
|
||||
remaining_steps: Remaining steps before recursion limit
|
||||
deployment_result: Final deployment result summary
|
||||
"""
|
||||
messages: Annotated[list[AnyMessage], operator.add]
|
||||
llm_calls: int
|
||||
remaining_steps: RemainingSteps
|
||||
deployment_result: dict[str, Any] | None
|
||||
|
||||
|
||||
# LLM call node
|
||||
def llm_call(state: dict):
|
||||
"""LLM decides whether to call a tool or respond"""
|
||||
messages = state.get("messages", [])
|
||||
|
||||
# Skip if no messages
|
||||
if not messages:
|
||||
return {"messages": []}
|
||||
|
||||
# Get model with tools
|
||||
model = create_base_model_with_tools(tools)
|
||||
|
||||
# Invoke model
|
||||
response = model.invoke(messages)
|
||||
|
||||
# Increment LLM call counter
|
||||
llm_calls = state.get("llm_calls", 0) + 1
|
||||
logger.info(f"LLM call #{llm_calls}: {type(response).__name__}")
|
||||
|
||||
return {
|
||||
"messages": [response],
|
||||
"llm_calls": llm_calls,
|
||||
}
|
||||
|
||||
|
||||
# Tool execution node
|
||||
def tool_execute(state: dict):
|
||||
"""Performs the tool call"""
|
||||
|
||||
result = []
|
||||
for tool_call in state["messages"][-1].tool_calls:
|
||||
tool = tools_by_name[tool_call["name"]]
|
||||
observation = tool.invoke(tool_call["args"])
|
||||
result.append(ToolMessage(
|
||||
content=observation,
|
||||
tool_call_id=tool_call["id"],
|
||||
name=tool_call["name"]
|
||||
))
|
||||
return {"messages": result}
|
||||
|
||||
|
||||
# Routing function
|
||||
def should_continue(state: dict):
|
||||
"""Decide whether to continue or end"""
|
||||
messages = state.get("messages", [])
|
||||
|
||||
if not messages:
|
||||
return END
|
||||
|
||||
last_message = messages[-1]
|
||||
|
||||
# Continue if last message has tool calls
|
||||
if isinstance(last_message, AIMessage) and last_message.tool_calls:
|
||||
return "tool_execute"
|
||||
|
||||
# End if last message is a regular text response
|
||||
return END
|
||||
|
||||
|
||||
# Build the graph
|
||||
def build_experiment_deploy_agent():
|
||||
"""Build the experiment deployment agent graph"""
|
||||
|
||||
# Create state graph
|
||||
workflow = StateGraph(ExperimentDeployState)
|
||||
|
||||
# Add nodes
|
||||
workflow.add_node("llm_call", llm_call)
|
||||
workflow.add_node("tool_execute", tool_execute)
|
||||
|
||||
# Set entry point
|
||||
workflow.add_edge(START, "llm_call")
|
||||
|
||||
# Add conditional edges
|
||||
workflow.add_conditional_edges(
|
||||
"llm_call",
|
||||
should_continue,
|
||||
{
|
||||
"tool_execute": "tool_execute",
|
||||
END: END,
|
||||
},
|
||||
)
|
||||
|
||||
# Add edge back to LLM after tool execution
|
||||
workflow.add_edge("tool_execute", "llm_call")
|
||||
|
||||
# Compile without checkpointer (experiment deploy is stateless)
|
||||
graph = workflow.compile()
|
||||
|
||||
logger.info("Experiment Deploy Agent graph compiled successfully")
|
||||
return graph
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_experiment_deploy_agent_graph = None
|
||||
|
||||
|
||||
def get_experiment_deploy_agent():
|
||||
"""Get or create the experiment deploy agent graph"""
|
||||
global _experiment_deploy_agent_graph
|
||||
if _experiment_deploy_agent_graph is None:
|
||||
_experiment_deploy_agent_graph = build_experiment_deploy_agent()
|
||||
return _experiment_deploy_agent_graph
|
||||
|
||||
|
||||
async def deploy_experiment_from_plan(
|
||||
plan_content: str,
|
||||
) -> Any:
|
||||
"""
|
||||
Deploy an experiment from plan content using the agent.
|
||||
|
||||
Args:
|
||||
plan_content: The experiment plan content
|
||||
|
||||
Yields:
|
||||
Agent execution updates
|
||||
"""
|
||||
agent = get_experiment_deploy_agent()
|
||||
|
||||
# Create system message
|
||||
system_msg = SystemMessage(content=EXPERIMENT_DEPLOY_SYSTEM_PROMPT)
|
||||
|
||||
# Create human message with the plan
|
||||
human_msg = HumanMessage(content=f"""
|
||||
Please deploy the following GNS3 experiment:
|
||||
|
||||
{plan_content}
|
||||
|
||||
Follow the deployment workflow and use the available tools to create the complete lab environment.
|
||||
""")
|
||||
|
||||
# Initial state
|
||||
initial_state = {
|
||||
"messages": [system_msg, human_msg],
|
||||
"llm_calls": 0,
|
||||
"deployment_result": None,
|
||||
}
|
||||
|
||||
# Stream agent execution (no config needed, stateless)
|
||||
async for event in agent.astream(initial_state):
|
||||
yield event
|
||||
|
||||
# Extract and yield deployment updates
|
||||
if "messages" in event:
|
||||
messages = event["messages"]
|
||||
for msg in messages:
|
||||
if isinstance(msg, ToolMessage):
|
||||
# Parse tool result and yield as update
|
||||
try:
|
||||
import json
|
||||
result = json.loads(msg.content)
|
||||
yield {"type": "tool_result", "tool": msg.name, "result": result}
|
||||
except:
|
||||
pass
|
||||
@ -1,415 +0,0 @@
|
||||
"""
|
||||
Qwen-VL Vision Model for Network Topology Recognition
|
||||
|
||||
This module provides vision recognition capabilities using Qwen-VL model
|
||||
through DashScope SDK to identify network topology diagrams from images.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import dashscope
|
||||
except ImportError:
|
||||
dashscope = None
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# System prompt for network topology recognition
|
||||
TOPOLOGY_RECOGNITION_PROMPT = """You are a professional network topology analysis expert. Please carefully analyze this network topology diagram and return the detailed information in JSON format.
|
||||
|
||||
Requirements:
|
||||
1. Identify all network devices (routers, switches, hosts, servers, etc.)
|
||||
2. Identify the connections between devices
|
||||
3. Identify interface and IP address information (if visible)
|
||||
4. Return in standard JSON format
|
||||
|
||||
Please return the result in the following JSON format:
|
||||
```json
|
||||
{
|
||||
"topology_name": "Topology name (inferred from content)",
|
||||
"description": "Brief description of the topology",
|
||||
"devices": [
|
||||
{
|
||||
"id": "unique device identifier",
|
||||
"name": "device name",
|
||||
"type": "device type (router/switch/host/server/cloud/firewall, etc.)",
|
||||
"model": "device model (if visible)",
|
||||
"position": {
|
||||
"x": 0,
|
||||
"y": 0
|
||||
}
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
{
|
||||
"id": "unique connection identifier",
|
||||
"source_device": "source device name",
|
||||
"source_interface": "source interface name (if visible)",
|
||||
"target_device": "target device name",
|
||||
"target_interface": "target interface name (if visible)",
|
||||
"link_type": "connection type (ethernet/serial, etc.)"
|
||||
}
|
||||
],
|
||||
"interfaces": [
|
||||
{
|
||||
"device": "device name",
|
||||
"interface": "interface name",
|
||||
"ip_address": "IP address (if visible)",
|
||||
"subnet_mask": "subnet mask (if visible)"
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"total_devices": 0,
|
||||
"total_links": 0,
|
||||
"device_types": {
|
||||
"router": 0,
|
||||
"switch": 0,
|
||||
"host": 0,
|
||||
"other": 0
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Notes:
|
||||
- If some information is not visible in the image, use null or empty string
|
||||
- Use the labels shown in the image for device names
|
||||
- Use relative coordinates for position (0-100 range)
|
||||
- Ensure valid JSON format, do not include markdown code block markers
|
||||
- Only return JSON, no other text explanations
|
||||
"""
|
||||
|
||||
|
||||
def encode_image_to_base64(image_path: str) -> str:
|
||||
"""
|
||||
Encode an image file to base64 string.
|
||||
|
||||
Args:
|
||||
image_path: Path to the image file
|
||||
|
||||
Returns:
|
||||
Base64 encoded string
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If image file doesn't exist
|
||||
IOError: If image file cannot be read
|
||||
"""
|
||||
path = Path(image_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
|
||||
try:
|
||||
with open(path, "rb") as image_file:
|
||||
base64_string = base64.b64encode(image_file.read()).decode("utf-8")
|
||||
logger.info(f"Successfully encoded image to base64: {image_path}")
|
||||
return base64_string
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to encode image: {e}")
|
||||
raise IOError(f"Failed to encode image {image_path}: {e}")
|
||||
|
||||
|
||||
def create_data_url(base64_string: str, mime_type: str = "image/png") -> str:
|
||||
"""
|
||||
Create a data URL from base64 string.
|
||||
|
||||
Args:
|
||||
base64_string: Base64 encoded image string
|
||||
mime_type: MIME type of the image (default: image/png)
|
||||
|
||||
Returns:
|
||||
Data URL string in format: data:[mime_type];base64,[base64_string]
|
||||
"""
|
||||
return f"data:{mime_type};base64,{base64_string}"
|
||||
|
||||
|
||||
class QwenVisionModel:
|
||||
"""
|
||||
Qwen-VL vision model wrapper for network topology recognition.
|
||||
|
||||
This class provides an interface to the Qwen-VL model through DashScope SDK
|
||||
for recognizing network topology diagrams from images.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
model_name: str = "qwen-vl-max",
|
||||
):
|
||||
"""
|
||||
Initialize Qwen-VL vision model.
|
||||
|
||||
Args:
|
||||
api_key: DashScope API key (if None, will load from config)
|
||||
model_name: Model name to use (default: qwen-vl-max)
|
||||
Options: qwen-vl-max, qwen-vl-plus, qwen-vl-v1
|
||||
|
||||
Raises:
|
||||
ImportError: If dashscope is not installed
|
||||
ValueError: If API key is not provided or found in config
|
||||
"""
|
||||
if dashscope is None:
|
||||
raise ImportError(
|
||||
"dashscope package is not installed. "
|
||||
"Please install it: pip install dashscope"
|
||||
)
|
||||
|
||||
# Load API key from environment variable if not provided
|
||||
if api_key is None:
|
||||
api_key = os.getenv("QWEN_API_KEY", "")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Qwen API key is required. Please set QWEN_API_KEY environment variable "
|
||||
"or pass it to the constructor."
|
||||
)
|
||||
|
||||
self.api_key = api_key
|
||||
self.model_name = model_name
|
||||
logger.info(f"QwenVisionModel initialized with model: {model_name}")
|
||||
|
||||
def _supports_streaming(self) -> bool:
|
||||
"""
|
||||
Check if the current model supports streaming output.
|
||||
|
||||
Returns:
|
||||
True if model supports streaming (qwen3-vl-plus, qwen3-vl-flash), False otherwise
|
||||
"""
|
||||
streaming_models = ["qwen3-vl-plus", "qwen3-vl-flash"]
|
||||
return self.model_name in streaming_models
|
||||
|
||||
def recognize_topology_from_file(
|
||||
self,
|
||||
image_path: str,
|
||||
prompt: str = TOPOLOGY_RECOGNITION_PROMPT,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Recognize network topology from an image file.
|
||||
|
||||
Args:
|
||||
image_path: Path to the image file
|
||||
prompt: Custom prompt for recognition (uses default if not provided)
|
||||
|
||||
Returns:
|
||||
Dictionary containing topology information
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If image file doesn't exist
|
||||
RuntimeError: If recognition fails
|
||||
"""
|
||||
logger.info(f"Recognizing topology from file: {image_path}")
|
||||
|
||||
# Encode image to base64
|
||||
base64_string = encode_image_to_base64(image_path)
|
||||
|
||||
# Detect MIME type from file extension
|
||||
mime_type = self._get_mime_type(image_path)
|
||||
|
||||
# Create data URL
|
||||
image_url = create_data_url(base64_string, mime_type)
|
||||
|
||||
# Call recognition
|
||||
return self.recognize_topology_from_base64(image_url, prompt)
|
||||
|
||||
def recognize_topology_from_base64(
|
||||
self,
|
||||
image_base64_or_url: str,
|
||||
prompt: str = TOPOLOGY_RECOGNITION_PROMPT,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Recognize network topology from a base64 encoded image or data URL.
|
||||
|
||||
Args:
|
||||
image_base64_or_url: Base64 encoded image string or data URL
|
||||
(format: data:image/xxx;base64,xxx)
|
||||
prompt: Custom prompt for recognition (uses default if not provided)
|
||||
|
||||
Returns:
|
||||
Dictionary containing topology information
|
||||
|
||||
Raises:
|
||||
RuntimeError: If recognition fails
|
||||
json.JSONDecodeError: If response is not valid JSON
|
||||
"""
|
||||
logger.info("Recognizing topology from base64 image")
|
||||
|
||||
# Use default prompt if not provided
|
||||
if prompt is None:
|
||||
prompt = TOPOLOGY_RECOGNITION_PROMPT
|
||||
|
||||
# Prepare message for DashScope (官方格式: 使用image字段传递Data URL)
|
||||
simple_prompt = """Analyze this network topology diagram and return the result in JSON format with the following structure:
|
||||
{
|
||||
"topology_name": "name",
|
||||
"description": "description",
|
||||
"devices": [{"id": "unique", "name": "device name", "type": "router/switch/host/etc", "model": "model if visible", "position": {"x": 0, "y": 0}}],
|
||||
"links": [{"id": "unique", "source_device": "name", "source_interface": "interface", "target_device": "name", "target_interface": "interface", "link_type": "ethernet/serial"}],
|
||||
"interfaces": [{"device": "name", "interface": "interface", "ip_address": "ip", "subnet_mask": "mask"}],
|
||||
"summary": {"total_devices": 0, "total_links": 0, "device_types": {"router": 0, "switch": 0, "host": 0, "other": 0}}
|
||||
}
|
||||
Return only valid JSON, no markdown code blocks."""
|
||||
|
||||
# image_base64_or_url should already be in data URL format (data:image/xxx;base64,xxx)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"image": image_base64_or_url},
|
||||
{"text": simple_prompt}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
logger.info(f"Model: {self.model_name}, Streaming: {self._supports_streaming()}")
|
||||
|
||||
try:
|
||||
# Determine if we should use streaming
|
||||
use_stream = self._supports_streaming()
|
||||
|
||||
if use_stream:
|
||||
# Call DashScope API with streaming
|
||||
logger.info("Using streaming mode for faster response")
|
||||
full_response = dashscope.MultiModalConversation.call(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
api_key=self.api_key,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect streaming response
|
||||
content = ""
|
||||
for chunk in full_response:
|
||||
# Extract content from streaming chunk
|
||||
if hasattr(chunk, 'output') and chunk.output:
|
||||
if hasattr(chunk.output, 'choices') and chunk.output.choices:
|
||||
choice = chunk.output.choices[0]
|
||||
if hasattr(choice, 'message'):
|
||||
message = choice.message
|
||||
# Handle different message formats in streaming
|
||||
if hasattr(message, 'content'):
|
||||
if message.content and len(message.content) > 0:
|
||||
content_item = message.content[0]
|
||||
if isinstance(content_item, dict):
|
||||
text = content_item.get('text', '')
|
||||
elif hasattr(content_item, 'text'):
|
||||
text = content_item.text
|
||||
else:
|
||||
text = str(content_item)
|
||||
content += text
|
||||
|
||||
logger.info(f"Completed streaming, total content length: {len(content)} chars")
|
||||
else:
|
||||
# Call DashScope API without streaming
|
||||
logger.info("Using non-streaming mode")
|
||||
response = dashscope.MultiModalConversation.call(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
api_key=self.api_key,
|
||||
)
|
||||
|
||||
# Extract content from non-streaming response
|
||||
try:
|
||||
if hasattr(response, 'output') and response.output:
|
||||
if hasattr(response.output, 'choices') and response.output.choices:
|
||||
if len(response.output.choices) > 0:
|
||||
choice = response.output.choices[0]
|
||||
if hasattr(choice, 'message') and choice.message:
|
||||
if hasattr(choice.message, 'content') and choice.message.content:
|
||||
if isinstance(choice.message.content, list) and len(choice.message.content) > 0:
|
||||
content_item = choice.message.content[0]
|
||||
if isinstance(content_item, dict):
|
||||
content = content_item.get('text', str(content_item))
|
||||
elif hasattr(content_item, 'text'):
|
||||
content = content_item.text
|
||||
else:
|
||||
content = str(content_item)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to extract content: {e}")
|
||||
content = ""
|
||||
|
||||
# Parse JSON response
|
||||
# Clean up any markdown code blocks
|
||||
content = content.strip()
|
||||
if content.startswith("```json"):
|
||||
content = content[7:]
|
||||
if content.startswith("```"):
|
||||
content = content[3:]
|
||||
if content.endswith("```"):
|
||||
content = content[:-3]
|
||||
content = content.strip()
|
||||
|
||||
try:
|
||||
topology_data = json.loads(content)
|
||||
logger.info("Successfully parsed topology JSON")
|
||||
return topology_data
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse JSON response: {e}")
|
||||
logger.error(f"Response content: {content[:500]}...")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to recognize topology: {e}")
|
||||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
raise RuntimeError(f"Failed to recognize topology: {e}") from e
|
||||
|
||||
def _get_mime_type(self, file_path: str) -> str:
|
||||
"""
|
||||
Get MIME type from file extension.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
MIME type string
|
||||
"""
|
||||
ext = Path(file_path).suffix.lower()
|
||||
mime_types = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".gif": "image/gif",
|
||||
".bmp": "image/bmp",
|
||||
".webp": "image/webp",
|
||||
}
|
||||
return mime_types.get(ext, "image/png")
|
||||
|
||||
|
||||
def create_qwen_vision_model(
|
||||
api_key: Optional[str] = None,
|
||||
model_name: Optional[str] = None,
|
||||
) -> QwenVisionModel:
|
||||
"""
|
||||
Factory function to create a Qwen-VL vision model instance.
|
||||
|
||||
This function loads configuration from environment variables and creates
|
||||
a QwenVisionModel instance with appropriate settings.
|
||||
|
||||
Args:
|
||||
api_key: DashScope API key (if None, will load from QWEN_API_KEY env var)
|
||||
model_name: Model name to use (if None, will load from QWEN_MODEL_NAME env var, default: qwen3-vl-plus)
|
||||
|
||||
Returns:
|
||||
QwenVisionModel instance
|
||||
|
||||
Raises:
|
||||
ImportError: If dashscope is not installed
|
||||
ValueError: If API key is not provided or found in environment
|
||||
|
||||
Example:
|
||||
>>> model = create_qwen_vision_model()
|
||||
>>> topology = model.recognize_topology_from_file("topology.png")
|
||||
>>> print(topology["topology_name"])
|
||||
>>> print(f"Found {len(topology['devices'])} devices")
|
||||
"""
|
||||
# Load model name from environment variable if not provided
|
||||
if model_name is None:
|
||||
model_name = os.getenv("QWEN_MODEL_NAME", "qwen3-vl-plus")
|
||||
|
||||
return QwenVisionModel(api_key=api_key, model_name=model_name)
|
||||
Loading…
x
Reference in New Issue
Block a user