Skip to content

Feat/all in one #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jlee-kitware wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Feat/all in one #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5446a13
Add MCP tool-call support for local and text-emitting backends
jlee-kitware Aug 1, 2026
f95569f
Support concurrent conversations, each owning its own state
jlee-kitware Aug 1, 2026
1708778
Add a console panel and data-file suggestions
jlee-kitware Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions src/vtk_prompt/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import click

from . import get_logger
from .client import VTKPromptClient
from .client import VTKPromptClient, load_conversation, save_conversation
from .provider_utils import DEFAULT_MODEL, DEFAULT_PROVIDER, get_default_model, supports_temperature

logger = get_logger(__name__)
Expand Down Expand Up @@ -138,13 +138,12 @@ def main(
temperature = 1.0

try:
client = VTKPromptClient(
verbose=verbose,
conversation_file=conversation,
mcp_url=mcp_url,
)
client = VTKPromptClient(verbose=verbose, mcp_url=mcp_url)
# The caller owns the conversation: load it, hand it to query, save it back.
messages = load_conversation(conversation)
result = client.query(
input_string,
conversation=messages,
api_key=token,
model=model,
base_url=base_url,
Expand All @@ -155,6 +154,7 @@ def main(
provider=provider,
custom_prompt=custom_prompt_data,
)
save_conversation(conversation, messages)

# Handle result with optional validation warnings
if isinstance(result, tuple):
Expand Down
227 changes: 143 additions & 84 deletions src/vtk_prompt/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,73 +33,79 @@
logger = get_logger(__name__)


@dataclass
class VTKPromptClient:
"""OpenAI client for VTK code generation."""

_instance: "VTKPromptClient" | None = None
_initialized: bool = False
verbose: bool = False
conversation_file: str | None = None
conversation: list[dict[str, str]] | None = None
mcp_url: str | None = None

def __new__(cls, **kwargs: Any) -> "VTKPromptClient":
"""Create singleton instance of VTKPromptClient."""
# Make sure that this is a singleton
if cls._instance is None:
cls._instance = super(VTKPromptClient, cls).__new__(cls)
cls._instance._initialized = False
cls._instance.conversation = []
return cls._instance

def __post_init__(self) -> None:
"""Post-init hook to prevent double initialization in singleton."""
if hasattr(self, "_initialized") and self._initialized:
return
self._initialized = True

def load_conversation(self) -> list[dict[str, str]]:
"""Load conversation history from file."""
if not self.conversation_file or not Path(self.conversation_file).exists():
return []

def _parse_text_tool_calls(content: str | None, tool_names: set[str]) -> list[dict] | None:
"""Extract tool calls a backend emitted as text instead of structured tool_calls.

Some local OpenAI-compatible backends (e.g. Ollama with a quantized model) return
a tool call as plain content rather than populating ``tool_calls``. Handle both a
``<tool_call>{...}</tool_call>`` block and a bare JSON object with ``name`` and
``arguments``. Only objects whose name matches a known tool are treated as calls,
so normal tagged answers are never misread. Returns a list of {name, arguments}.
"""
if not content:
return None
candidates = re.findall(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", content, re.DOTALL)
if not candidates:
stripped = content.strip()
if stripped.startswith("{") and stripped.endswith("}"):
candidates = [stripped]
calls: list[dict] = []
for raw in candidates:
try:
with open(self.conversation_file, "r") as f:
data = json.load(f)
if isinstance(data, list):
return data
else:
logger.warning("Invalid conversation file format, no history loaded.")
return []
except Exception as e:
logger.error("Could not load conversation file: %s", e)
obj = json.loads(raw)
except (ValueError, TypeError):
continue
name = obj.get("name")
args = obj.get("arguments", {})
if isinstance(args, str):
try:
args = json.loads(args)
except (ValueError, TypeError):
args = {}
if name in tool_names and isinstance(args, dict):
calls.append({"name": name, "arguments": args})
return calls or None


def load_conversation(path: str | None) -> list[dict[str, str]]:
"""Load conversation history from a file path."""
if not path or not Path(path).exists():
return []
try:
with open(path, "r") as f:
data = json.load(f)
if isinstance(data, list):
return data
logger.warning("Invalid conversation file format, no history loaded.")
return []
except Exception as e:
logger.error("Could not load conversation file: %s", e)
return []

def save_conversation(self) -> None:
"""Save conversation history to file."""
if not self.conversation_file or not self.conversation:
return

try:
# Ensure directory exists
Path(self.conversation_file).parent.mkdir(parents=True, exist_ok=True)
def save_conversation(path: str | None, messages: list[dict[str, str]]) -> None:
"""Save conversation history to a file path."""
if not path or not messages:
return
try:
Path(path).parent.mkdir(parents=True, exist_ok=True)
with open(path, "w") as f:
json.dump(messages, f, indent=2)
except Exception as e:
logger.error("Could not save conversation file: %s", e)

with open(self.conversation_file, "w") as f:
json.dump(self.conversation, f, indent=2)
except Exception as e:
logger.error("Could not save conversation file: %s", e)

def update_conversation(
self, new_convo: list[dict[str, str]], new_convo_file: str | None = None
) -> None:
"""Update conversation history with new conversation."""
if not self.conversation:
self.conversation = []
self.conversation.extend(new_convo)
@dataclass
class VTKPromptClient:
"""Stateless OpenAI client for VTK code generation.

if new_convo_file:
self.conversation_file = new_convo_file
The client holds only transport configuration. Conversation state belongs to
the caller (a UI session or the CLI), which passes its message list into
query() and owns it afterwards.
"""

verbose: bool = False
mcp_url: str | None = None

def validate_code_syntax(self, code_string: str) -> tuple[bool, str | None]:
"""Validate Python code syntax using AST."""
Expand Down Expand Up @@ -296,6 +302,9 @@ def query(
custom_prompt: dict | None = None,
ui_mode: bool = False,
execution_error: str | None = None,
log_tool_calls: bool = False,
agentic_retrieval: bool = False,
conversation: list[dict[str, str]] | None = None,
) -> tuple[str, str, Any] | tuple[str, str, Any, list[str]] | str:
"""Generate VTK code using vtk-mcp tools when available.

Expand All @@ -321,11 +330,11 @@ def query(
# Create client with current parameters
client = openai.OpenAI(api_key=api_key, base_url=base_url)

# Load existing conversation if present
if self.conversation_file and not self.conversation:
self.conversation = self.load_conversation()
# The caller owns the conversation; append to it in place so the result is
# visible to whoever passed it in. No conversation state lives on self.
messages: list[dict[str, str]] = conversation if conversation is not None else []

if not message and not self.conversation:
if not message and not messages:
raise ValueError("No prompt or conversation file provided")

# Set up vtk-mcp client (context retrieval, tool calling, and code validation)
Expand All @@ -343,9 +352,9 @@ def query(

if execution_error:
# Retry after execution failure: append error and let LLM fix it with tools
if not self.conversation:
if not messages:
raise ValueError("No conversation to retry")
self.conversation.append(
messages.append(
{
"role": "user",
"content": (
Expand All @@ -359,7 +368,8 @@ def query(
else:
# Normal path: build context and prompt
context_snippets = None
if mcp_client:
# Agentic mode: skip pre-injected context so the model must use tools.
if mcp_client and not agentic_retrieval:
mcp_context = mcp_client.get_enriched_context(message, top_k=top_k)
if mcp_context:
context_snippets = mcp_context
Expand Down Expand Up @@ -410,13 +420,18 @@ def query(
mcp_str = " + MCP" if mcp_client else ""
logger.debug(f"Using component assembly ({mode_str}{mcp_str})")

if not self.conversation:
self.conversation = list(yaml_messages)
if not messages:
messages.extend(yaml_messages)
elif message and yaml_messages:
self.conversation.append(yaml_messages[-1])
messages.append(yaml_messages[-1])

# Fetch vtk-mcp tools for LLM tool calling
tools = mcp_client.list_tools() if mcp_client else []
tool_names: set[str] = {
str(name)
for t in tools
if (name := (t.get("function") or {}).get("name")) is not None
}

# Retry loop for AST validation
for attempt in range(retry_attempts):
Expand All @@ -432,7 +447,7 @@ def query(
for _ in range(MAX_TOOL_ROUNDS):
response = client.chat.completions.create(
model=model,
messages=self.conversation, # type: ignore[arg-type]
messages=messages, # type: ignore[arg-type]
max_completion_tokens=max_tokens,
temperature=temperature,
**( # type: ignore[call-overload]
Expand All @@ -459,7 +474,7 @@ def query(
}
for tc in choice.message.tool_calls
]
self.conversation.append(tc_msg)
messages.append(tc_msg)

# Execute each tool and append results
for tc in choice.message.tool_calls:
Expand All @@ -468,18 +483,64 @@ def query(
except Exception:
args = {}
result = mcp_client.call_tool(tc.function.name, args) # type: ignore
logger.debug("Tool %s -> %s...", tc.function.name, result[:80])
self.conversation.append(
if log_tool_calls:
logger.info(
"vtk-mcp: %s(%s) -> %s",
tc.function.name,
tc.function.arguments,
result[:120],
)
else:
logger.debug("Tool %s -> %s...", tc.function.name, result[:80])
messages.append(
{"role": "tool", "tool_call_id": tc.id, "content": result}
)
continue # let LLM decide what to do next

# Fallback: the backend returned a tool call as plain text rather than
# in tool_calls (common with quantized local models). Run it anyway.
text_calls = _parse_text_tool_calls(choice.message.content, tool_names)
if tools and text_calls:
text_msg: dict = {"role": "assistant", "content": ""}
text_msg["tool_calls"] = [
{
"id": f"call_{i}",
"type": "function",
"function": {
"name": c["name"],
"arguments": json.dumps(c["arguments"]),
},
}
for i, c in enumerate(text_calls)
]
messages.append(text_msg)
for i, c in enumerate(text_calls):
result = mcp_client.call_tool(c["name"], c["arguments"]) # type: ignore
if log_tool_calls:
logger.info(
"vtk-mcp (text): %s(%s) -> %s",
c["name"],
json.dumps(c["arguments"]),
result[:120],
)
else:
logger.debug("Tool %s (text) -> %s...", c["name"], result[:80])
messages.append(
{"role": "tool", "tool_call_id": f"call_{i}", "content": result}
)
continue

# LLM generated a response (not a tool call)
content = choice.message.content or "No content in response"
break

if content is None:
# Tool loop exhausted without a text response
if log_tool_calls:
logger.info(
"vtk-mcp: tool loop hit the %d-round cap without a final response",
MAX_TOOL_ROUNDS,
)
if attempt == retry_attempts - 1:
return ("No response generated", "", getattr(response, "usage", None) or {})
continue
Expand All @@ -500,8 +561,8 @@ def query(

if not expl_matches or not code_matches:
if attempt < retry_attempts - 1:
self.conversation.append({"role": "assistant", "content": content})
self.conversation.append(
messages.append({"role": "assistant", "content": content})
messages.append(
{
"role": "user",
"content": (
Expand All @@ -526,8 +587,8 @@ def query(
if vtk_diagnostics:
if attempt < retry_attempts - 1:
logger.debug("vtk-mcp validation issues found, retrying")
self.conversation.append({"role": "assistant", "content": content})
self.conversation.append(
messages.append({"role": "assistant", "content": content})
messages.append(
{
"role": "user",
"content": (
Expand All @@ -543,8 +604,7 @@ def query(
f"VTK API issues found:\n{vtk_diagnostics}"
)
if message:
self.conversation.append({"role": "assistant", "content": content})
self.save_conversation()
messages.append({"role": "assistant", "content": content})
if validation_warnings:
return (
generated_explanation,
Expand All @@ -557,8 +617,8 @@ def query(
elif attempt < retry_attempts - 1:
if self.verbose:
logger.warning("AST validation failed: %s. Retrying...", error_msg)
self.conversation.append({"role": "assistant", "content": content})
self.conversation.append(
messages.append({"role": "assistant", "content": content})
messages.append(
{
"role": "user",
"content": (
Expand All @@ -571,8 +631,7 @@ def query(
if self.verbose:
logger.error("Final attempt failed AST validation: %s", error_msg)
if message:
self.conversation.append({"role": "assistant", "content": content})
self.save_conversation()
messages.append({"role": "assistant", "content": content})
if validation_warnings:
return (
generated_explanation,
Expand Down
Loading
Loading