From 4522412be19e748b266d5fc43939fed6e0fd7471 Mon Sep 17 00:00:00 2001 From: Kenneth Belitzky Date: Sun, 29 Jun 2025 03:25:57 +0000 Subject: [PATCH 1/3] Enhance environment configuration and implement model wrapper for AI content generation --- .devcontainer/devcontainer.json | 1 - .env.example | 28 ++++++++++++++++++++ requirements.txt | 1 + struct_module/file_item.py | 45 +++++---------------------------- struct_module/model_wrapper.py | 32 +++++++++++++++++++++++ 5 files changed, 68 insertions(+), 39 deletions(-) create mode 100644 struct_module/model_wrapper.py diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 98f2d9f..b413378 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -2,7 +2,6 @@ "name": "Struct devcontainer", "image": "mcr.microsoft.com/devcontainers/python:3", "features": { - "ghcr.io/devcontainers/features/python:1": {}, "ghcr.io/gvatsal60/dev-container-features/pre-commit": {}, "ghcr.io/eitsupi/devcontainer-features/go-task:latest": {} }, diff --git a/.env.example b/.env.example index f0ac698..a90fed4 100644 --- a/.env.example +++ b/.env.example @@ -1 +1,29 @@ +# The AI model to use, e.g., "openai:gpt-4.1" +# check for available models at https://ai.pydantic.dev/models/ +AI_MODEL=your-model-name-here + +# To use openai's API, you need to set your API key in the environment variable OPENAI_API_KEY. OPENAI_API_KEY=your-api-key-here + +# To use Anthropy's API, you need to set your API key in the environment variable ANTHROPY_API_KEY. +ANTHROPY_API_KEY=your-anthropy-api-key-here + +# To use Gemini's API, you need to set your API key in the environment variable GEMINI_API_KEY. +GEMINI_API_KEY=your-gemini-api-key-here + +# To use Google's API, you need to set your API key in the environment variable GOOGLE_API_KEY. +GOOGLE_API_KEY=your-google-api-key-here + +# Bedrock API key +AWS_ACCESS_KEY_ID=your-aws-access-key-id-here +AWS_SECRET_ACCESS_KEY=your-aws-secret-access-key-here +AWS_DEFAULT_REGION=your-aws-region-here + +# To use Cohere's API, you need to set your API key in the environment variable CO_API_KEY. +CO_API_KEY=your-cohere-api-key-here + +# To use Groq's API, you need to set your API key in the environment variable GROQ_API_KEY. +GROQ_API_KEY=your-groq-api-key-here + +# To use Mistral's API, you need to set your API key in the environment variable MISTRAL_API_KEY. +MISTRAL_API_KEY=your-mistral-api-key-here diff --git a/requirements.txt b/requirements.txt index 8704ed1..6d65028 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,3 +10,4 @@ boto3 google-cloud google-api-core cachetools +pydantic-ai diff --git a/struct_module/file_item.py b/struct_module/file_item.py index e5a7fcd..d321f05 100644 --- a/struct_module/file_item.py +++ b/struct_module/file_item.py @@ -4,16 +4,13 @@ import shutil import logging import time -from openai import OpenAI from dotenv import load_dotenv from struct_module.template_renderer import TemplateRenderer from struct_module.content_fetcher import ContentFetcher +from struct_module.model_wrapper import ModelWrapper load_dotenv() -openai_api_key = os.getenv("OPENAI_API_KEY") -openai_model = os.getenv("OPENAI_MODEL") - class FileItem: def __init__(self, properties): self.logger = logging.getLogger(__name__) @@ -32,11 +29,9 @@ def __init__(self, properties): self.system_prompt = properties.get("system_prompt") or properties.get("global_system_prompt") self.user_prompt = properties.get("user_prompt") - self.openai_client = None self.mappings = properties.get("mappings", {}) - if openai_api_key: - self._configure_openai() + self.model_wrapper = ModelWrapper(self.logger) self.template_renderer = TemplateRenderer( self.config_variables, @@ -45,30 +40,17 @@ def __init__(self, properties): self.mappings ) - def _configure_openai(self): - self.openai_client = OpenAI(api_key=openai_api_key) - if not openai_model: - self.logger.debug("OpenAI model not found. Using default model.") - self.openai_model = "gpt-4.1" - else: - self.logger.debug(f"Using OpenAI model: {openai_model}") - self.openai_model = openai_model - def _get_file_directory(self): return os.path.dirname(self.name) def process_prompt(self, dry_run=False, existing_content=None): if self.user_prompt: - if not self.openai_client or not openai_api_key: - self.logger.warning("Skipping processing prompt as OpenAI API key is not set.") - return if not self.system_prompt: system_prompt = "You are a software developer working on a project. You need to create a file with the following content:" else: system_prompt = self.system_prompt - # If existing_content is provided, append it to the user prompt user_prompt = self.user_prompt if existing_content: user_prompt += f"\n\nCurrent file content (if any):\n```\n{existing_content}\n```\n\nPlease modify existing content so that it meets the new requirements. Your output should be plain text, without any code blocks or formatting. Do not include any explanations or comments. Just provide the final content of the file." @@ -76,24 +58,11 @@ def process_prompt(self, dry_run=False, existing_content=None): self.logger.debug(f"Using system prompt: {system_prompt}") self.logger.debug(f"Using user prompt: {user_prompt}") - if dry_run: - self.logger.info("[DRY RUN] Would generate content using OpenAI API.") - self.content = "[DRY RUN] Generating content using OpenAI" - return - - if self.openai_client and openai_api_key: - completion = self.openai_client.chat.completions.create( - model=self.openai_model, - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt} - ] - ) - - self.content = completion.choices[0].message.content - else: - self.content = "OpenAI API key not found. Skipping content generation." - self.logger.warning("Skipping processing prompt as OpenAI API key is not set.") + self.content = self.model_wrapper.generate_content( + system_prompt, + user_prompt, + dry_run=dry_run + ) self.logger.debug(f"Generated content: \n\n{self.content}") def fetch_content(self): diff --git a/struct_module/model_wrapper.py b/struct_module/model_wrapper.py new file mode 100644 index 0000000..782fc90 --- /dev/null +++ b/struct_module/model_wrapper.py @@ -0,0 +1,32 @@ +import os +import logging +from dotenv import load_dotenv +from pydantic_ai import Agent + +load_dotenv() + +class ModelWrapper: + """ + Wraps model logic using pydantic-ai Agent, allowing use of multiple LLM providers. + """ + def __init__(self, logger=None): + self.logger = logger or logging.getLogger(__name__) + self.model_name = os.getenv("AI_MODEL") or "openai:gpt-4.1" + self.agent = Agent(model=self.model_name) + self.logger.debug(f"Configured Agent with model: {self.model_name}") + + def generate_content(self, system_prompt, user_prompt, dry_run=False): + if not self.agent: + self.logger.warning("No agent configured. Skipping content generation.") + return "No agent configured. Skipping content generation." + if dry_run: + self.logger.info("[DRY RUN] Would generate content using AI agent.") + return "[DRY RUN] Generating content using AI agent" + prompt = f"{user_prompt}" + try: + self.agent.system_prompt = system_prompt + result = self.agent.run_sync(prompt) + return result.output + except Exception as e: + self.logger.error(f"AI agent generation failed: {e}") + return f"AI agent generation failed: {e}" From 2e59eadef50a06a87c6f4486ec4f6c50465c6f93 Mon Sep 17 00:00:00 2001 From: Kenneth Belitzky Date: Sun, 29 Jun 2025 15:55:43 +0000 Subject: [PATCH 2/3] Update devcontainer configuration to include pre-commit feature --- .devcontainer/devcontainer.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 8fbe21f..4764e1c 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -2,7 +2,8 @@ "name": "Struct devcontainer", "image": "mcr.microsoft.com/devcontainers/python:3", "features": { - "ghcr.io/gvatsal60/dev-container-features/pre-commit": {}, + "ghcr.io/devcontainers/features/python:1": {}, + "ghcr.io/gvatsal60/dev-container-features/pre-commit:1": {} }, "postCreateCommand": "bash ./scripts/devcontainer_start.sh", "customizations": { From 89baaeee642180014b664dd0ed6859f1a6208e9e Mon Sep 17 00:00:00 2001 From: Kenneth Belitzky Date: Mon, 30 Jun 2025 03:54:56 +0000 Subject: [PATCH 3/3] Add environment variable for OpenAI API key in test script --- .github/workflows/test-script.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/test-script.yaml b/.github/workflows/test-script.yaml index 5c7a033..605a542 100644 --- a/.github/workflows/test-script.yaml +++ b/.github/workflows/test-script.yaml @@ -8,6 +8,9 @@ on: branches: - main +env: + OPENAI_API_KEY: "my-test-key" + jobs: build: