janito 0.3.0__tar.gz → 0.5.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (66) hide show
  1. janito-0.5.0/PKG-INFO +146 -0
  2. janito-0.5.0/README.md +122 -0
  3. janito-0.5.0/janito/__init__.py +49 -0
  4. janito-0.5.0/janito/__main__.py +54 -0
  5. janito-0.5.0/janito/_contextparser.py +113 -0
  6. janito-0.5.0/janito/agents/__init__.py +22 -0
  7. janito-0.5.0/janito/agents/agent.py +21 -0
  8. janito-0.5.0/janito/agents/claudeai.py +64 -0
  9. janito-0.5.0/janito/agents/openai.py +53 -0
  10. janito-0.5.0/janito/agents/test.py +34 -0
  11. janito-0.5.0/janito/analysis/__init__.py +33 -0
  12. janito-0.5.0/janito/analysis/display.py +149 -0
  13. janito-0.5.0/janito/analysis/options.py +112 -0
  14. janito-0.5.0/janito/analysis/prompts.py +75 -0
  15. janito-0.5.0/janito/change/__init__.py +19 -0
  16. janito-0.5.0/janito/change/applier.py +269 -0
  17. janito-0.5.0/janito/change/content.py +62 -0
  18. janito-0.5.0/janito/change/indentation.py +33 -0
  19. janito-0.5.0/janito/change/position.py +169 -0
  20. janito-0.5.0/janito/changehistory.py +46 -0
  21. janito-0.5.0/janito/changeviewer/__init__.py +12 -0
  22. janito-0.5.0/janito/changeviewer/diff.py +28 -0
  23. janito-0.5.0/janito/changeviewer/panels.py +268 -0
  24. janito-0.5.0/janito/changeviewer/styling.py +59 -0
  25. janito-0.5.0/janito/changeviewer/themes.py +57 -0
  26. janito-0.5.0/janito/cli/__init__.py +2 -0
  27. janito-0.5.0/janito/cli/commands.py +53 -0
  28. janito-0.5.0/janito/cli/functions.py +286 -0
  29. janito-0.5.0/janito/cli/registry.py +26 -0
  30. janito-0.5.0/janito/common.py +23 -0
  31. {janito-0.3.0 → janito-0.5.0}/janito/config.py +8 -3
  32. janito-0.5.0/janito/console/__init__.py +3 -0
  33. janito-0.5.0/janito/console/commands.py +112 -0
  34. janito-0.5.0/janito/console/core.py +62 -0
  35. janito-0.5.0/janito/console/display.py +157 -0
  36. janito-0.5.0/janito/fileparser.py +334 -0
  37. janito-0.5.0/janito/prompts.py +81 -0
  38. janito-0.5.0/janito/qa.py +65 -0
  39. janito-0.5.0/janito/review.py +13 -0
  40. {janito-0.3.0 → janito-0.5.0}/janito/scan.py +68 -14
  41. janito-0.5.0/janito/tests/test_fileparser.py +26 -0
  42. janito-0.5.0/janito/version.py +23 -0
  43. {janito-0.3.0 → janito-0.5.0}/pyproject.toml +3 -2
  44. janito-0.5.0/tests/conftest.py +45 -0
  45. janito-0.5.0/tests/test_contentchange.py +68 -0
  46. janito-0.5.0/tests/test_integration.py +64 -0
  47. janito-0.5.0/tests/test_prompts.py +59 -0
  48. janito-0.5.0/tests/test_scan.py +98 -0
  49. janito-0.5.0/tools/release.sh +77 -0
  50. janito-0.3.0/CHANGELOG.md +0 -26
  51. janito-0.3.0/PKG-INFO +0 -138
  52. janito-0.3.0/README.md +0 -115
  53. janito-0.3.0/janito/__init__.py +0 -2
  54. janito-0.3.0/janito/__main__.py +0 -260
  55. janito-0.3.0/janito/changeviewer.py +0 -64
  56. janito-0.3.0/janito/claude.py +0 -74
  57. janito-0.3.0/janito/console.py +0 -60
  58. janito-0.3.0/janito/contentchange.py +0 -165
  59. janito-0.3.0/janito/prompts.py +0 -97
  60. janito-0.3.0/janito/qa.py +0 -32
  61. janito-0.3.0/tests/__init__.py +0 -2
  62. janito-0.3.0/tests/test_main.py +0 -51
  63. janito-0.3.0/tools/release.sh +0 -28
  64. {janito-0.3.0 → janito-0.5.0}/.gitignore +0 -0
  65. {janito-0.3.0 → janito-0.5.0}/LICENSE +0 -0
  66. {janito-0.3.0 → janito-0.5.0}/setup.py +0 -0
janito-0.5.0/PKG-INFO ADDED
@@ -0,0 +1,146 @@
1
+ Metadata-Version: 2.3
2
+ Name: janito
3
+ Version: 0.5.0
4
+ Summary: A CLI tool for software development tasks powered by AI
5
+ Project-URL: Homepage, https://github.com/joaompinto/janito
6
+ Project-URL: Repository, https://github.com/joaompinto/janito.git
7
+ Author-email: João Pinto <lamego.pinto@gmail.com>
8
+ License: MIT
9
+ Classifier: Development Status :: 4 - Beta
10
+ Classifier: Environment :: Console
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Programming Language :: Python :: 3.8
14
+ Classifier: Programming Language :: Python :: 3.9
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Topic :: Software Development
17
+ Requires-Python: >=3.8
18
+ Requires-Dist: anthropic
19
+ Requires-Dist: pathspec
20
+ Requires-Dist: rich
21
+ Requires-Dist: tomli
22
+ Requires-Dist: typer
23
+ Description-Content-Type: text/markdown
24
+
25
+ # 🤖 Janito CLI
26
+
27
+ A CLI tool for software development tasks powered by AI. Janito is your friendly AI-powered software development buddy that helps with coding tasks like refactoring, documentation updates, and code optimization.
28
+
29
+ ## 📥 Installation
30
+
31
+ 1. Install using pip:
32
+ ```bash
33
+ pip install janito
34
+ ```
35
+
36
+ 2. Verify installation:
37
+ ```bash
38
+ janito --version
39
+ ```
40
+
41
+ ## ⚙️ Setup
42
+
43
+ 1. Get your Anthropic API key from [Anthropic's website](https://www.anthropic.com/)
44
+
45
+ 2. Set your API key:
46
+ ```bash
47
+ # Linux/macOS
48
+ export ANTHROPIC_API_KEY='your-api-key-here'
49
+
50
+ # Windows (Command Prompt)
51
+ set ANTHROPIC_API_KEY=your-api-key-here
52
+
53
+ # Windows (PowerShell)
54
+ $env:ANTHROPIC_API_KEY='your-api-key-here'
55
+ ```
56
+
57
+ 3. (Optional) Configure default test command:
58
+ ```bash
59
+ export JANITO_TEST_CMD='pytest' # or your preferred test command
60
+ ```
61
+
62
+ ## 🚀 Quick Start
63
+
64
+ ### Basic Usage
65
+
66
+ ```bash
67
+ # Add docstrings to your code
68
+ janito "add docstrings to this file"
69
+
70
+ # Optimize a function
71
+ janito "optimize the main function"
72
+
73
+ # Get code explanations
74
+ janito --ask "explain this code"
75
+ ```
76
+
77
+ ### Common Scenarios
78
+
79
+ 1. **Code Refactoring**
80
+ ```bash
81
+ # Refactor with test validation
82
+ janito "refactor this code to use list comprehension" --test "pytest"
83
+
84
+ # Refactor specific directory
85
+ janito "update imports" -i ./src
86
+ ```
87
+
88
+ 2. **Documentation Updates**
89
+ ```bash
90
+ # Add or update docstrings
91
+ janito "add type hints and docstrings"
92
+
93
+ # Generate README
94
+ janito "create a README for this project"
95
+ ```
96
+
97
+ 3. **Code Analysis**
98
+ ```bash
99
+ # Get code explanations
100
+ janito --ask "what does this function do?"
101
+
102
+ # Find potential improvements
103
+ janito --ask "suggest optimizations for this code"
104
+ ```
105
+
106
+ ## 🛠️ Command Reference
107
+
108
+ ### Syntax
109
+ ```bash
110
+ janito [OPTIONS] [REQUEST]
111
+ ```
112
+
113
+ ### Key Options
114
+
115
+ | Option | Description |
116
+ |--------|-------------|
117
+ | `REQUEST` | The AI request/instruction (in quotes) |
118
+ | `-w, --working-dir PATH` | Working directory [default: current] |
119
+ | `-i, --include PATH` | Include directory int the working context (can be multiple)|
120
+ | `--ask QUESTION` | Ask questions without making changes |
121
+ | `--test COMMAND` | Run tests before applying changes |
122
+ | `--debug` | Enable debug logging |
123
+ | `--verbose` | Enable verbose mode |
124
+ | `--version` | Show version information |
125
+ | `--help` | Show help message |
126
+
127
+ ## 🔑 Key Features
128
+
129
+ - 🤖 AI-powered code analysis and modifications
130
+ - 💻 Interactive console mode
131
+ - ✅ Syntax validation for Python files
132
+ - 👀 Change preview and confirmation
133
+ - 🧪 Test command execution
134
+ - 📜 Change history tracking
135
+
136
+ ## 📚 Additional Information
137
+
138
+ - Requires Python 3.8+
139
+ - Changes are backed up in `.janito/changes_history/`
140
+ - Environment variables:
141
+ - `ANTHROPIC_API_KEY`: Required for API access
142
+ - `JANITO_TEST_CMD`: Default test command (optional)
143
+
144
+ ## 📄 License
145
+
146
+ MIT License - see [LICENSE](LICENSE) file for details.
janito-0.5.0/README.md ADDED
@@ -0,0 +1,122 @@
1
+ # 🤖 Janito CLI
2
+
3
+ A CLI tool for software development tasks powered by AI. Janito is your friendly AI-powered software development buddy that helps with coding tasks like refactoring, documentation updates, and code optimization.
4
+
5
+ ## 📥 Installation
6
+
7
+ 1. Install using pip:
8
+ ```bash
9
+ pip install janito
10
+ ```
11
+
12
+ 2. Verify installation:
13
+ ```bash
14
+ janito --version
15
+ ```
16
+
17
+ ## ⚙️ Setup
18
+
19
+ 1. Get your Anthropic API key from [Anthropic's website](https://www.anthropic.com/)
20
+
21
+ 2. Set your API key:
22
+ ```bash
23
+ # Linux/macOS
24
+ export ANTHROPIC_API_KEY='your-api-key-here'
25
+
26
+ # Windows (Command Prompt)
27
+ set ANTHROPIC_API_KEY=your-api-key-here
28
+
29
+ # Windows (PowerShell)
30
+ $env:ANTHROPIC_API_KEY='your-api-key-here'
31
+ ```
32
+
33
+ 3. (Optional) Configure default test command:
34
+ ```bash
35
+ export JANITO_TEST_CMD='pytest' # or your preferred test command
36
+ ```
37
+
38
+ ## 🚀 Quick Start
39
+
40
+ ### Basic Usage
41
+
42
+ ```bash
43
+ # Add docstrings to your code
44
+ janito "add docstrings to this file"
45
+
46
+ # Optimize a function
47
+ janito "optimize the main function"
48
+
49
+ # Get code explanations
50
+ janito --ask "explain this code"
51
+ ```
52
+
53
+ ### Common Scenarios
54
+
55
+ 1. **Code Refactoring**
56
+ ```bash
57
+ # Refactor with test validation
58
+ janito "refactor this code to use list comprehension" --test "pytest"
59
+
60
+ # Refactor specific directory
61
+ janito "update imports" -i ./src
62
+ ```
63
+
64
+ 2. **Documentation Updates**
65
+ ```bash
66
+ # Add or update docstrings
67
+ janito "add type hints and docstrings"
68
+
69
+ # Generate README
70
+ janito "create a README for this project"
71
+ ```
72
+
73
+ 3. **Code Analysis**
74
+ ```bash
75
+ # Get code explanations
76
+ janito --ask "what does this function do?"
77
+
78
+ # Find potential improvements
79
+ janito --ask "suggest optimizations for this code"
80
+ ```
81
+
82
+ ## 🛠️ Command Reference
83
+
84
+ ### Syntax
85
+ ```bash
86
+ janito [OPTIONS] [REQUEST]
87
+ ```
88
+
89
+ ### Key Options
90
+
91
+ | Option | Description |
92
+ |--------|-------------|
93
+ | `REQUEST` | The AI request/instruction (in quotes) |
94
+ | `-w, --working-dir PATH` | Working directory [default: current] |
95
+ | `-i, --include PATH` | Include directory int the working context (can be multiple)|
96
+ | `--ask QUESTION` | Ask questions without making changes |
97
+ | `--test COMMAND` | Run tests before applying changes |
98
+ | `--debug` | Enable debug logging |
99
+ | `--verbose` | Enable verbose mode |
100
+ | `--version` | Show version information |
101
+ | `--help` | Show help message |
102
+
103
+ ## 🔑 Key Features
104
+
105
+ - 🤖 AI-powered code analysis and modifications
106
+ - 💻 Interactive console mode
107
+ - ✅ Syntax validation for Python files
108
+ - 👀 Change preview and confirmation
109
+ - 🧪 Test command execution
110
+ - 📜 Change history tracking
111
+
112
+ ## 📚 Additional Information
113
+
114
+ - Requires Python 3.8+
115
+ - Changes are backed up in `.janito/changes_history/`
116
+ - Environment variables:
117
+ - `ANTHROPIC_API_KEY`: Required for API access
118
+ - `JANITO_TEST_CMD`: Default test command (optional)
119
+
120
+ ## 📄 License
121
+
122
+ MIT License - see [LICENSE](LICENSE) file for details.
@@ -0,0 +1,49 @@
1
+ """Core package initialization for Janito."""
2
+
3
+ from .analysis import (
4
+ AnalysisOption,
5
+ parse_analysis_options,
6
+ format_analysis,
7
+ get_history_file_type,
8
+ get_history_path,
9
+ get_timestamp,
10
+ save_to_file,
11
+ build_request_analysis_prompt,
12
+ get_option_selection,
13
+ prompt_user,
14
+ validate_option_letter
15
+ )
16
+
17
+ from .change import (
18
+ apply_single_change,
19
+ parse_and_apply_changes_sequence,
20
+ get_file_type,
21
+ process_and_save_changes,
22
+ format_parsed_changes,
23
+ apply_content_changes,
24
+ handle_changes_file
25
+ )
26
+
27
+ __all__ = [
28
+ # Analysis exports
29
+ 'AnalysisOption',
30
+ 'parse_analysis_options',
31
+ 'format_analysis',
32
+ 'get_history_file_type',
33
+ 'get_history_path',
34
+ 'get_timestamp',
35
+ 'save_to_file',
36
+ 'build_request_analysis_prompt',
37
+ 'get_option_selection',
38
+ 'prompt_user',
39
+ 'validate_option_letter',
40
+
41
+ # Change exports
42
+ 'apply_single_change',
43
+ 'parse_and_apply_changes_sequence',
44
+ 'get_file_type',
45
+ 'process_and_save_changes',
46
+ 'format_parsed_changes',
47
+ 'apply_content_changes',
48
+ 'handle_changes_file'
49
+ ]
@@ -0,0 +1,54 @@
1
+ import typer
2
+ from typing import Optional, List
3
+ from pathlib import Path
4
+ from rich.console import Console
5
+ from .version import get_version
6
+
7
+ from janito.agents import AgentSingleton
8
+ from janito.config import config
9
+
10
+ from .cli.commands import handle_request, handle_ask, handle_play, handle_scan
11
+
12
+ app = typer.Typer(add_completion=False)
13
+
14
+ def typer_main(
15
+ change_request: str = typer.Argument(None, help="Change request or command"),
16
+ workdir: Optional[Path] = typer.Option(None, "-w", "--workdir", help="Working directory", file_okay=False, dir_okay=True),
17
+ debug: bool = typer.Option(False, "--debug", help="Show debug information"),
18
+ verbose: bool = typer.Option(False, "--verbose", help="Show verbose output"),
19
+ include: Optional[List[Path]] = typer.Option(None, "-i", "--include", help="Additional paths to include"),
20
+ ask: Optional[str] = typer.Option(None, "--ask", help="Ask a question about the codebase"),
21
+ play: Optional[Path] = typer.Option(None, "--play", help="Replay a saved prompt file"),
22
+ version: bool = typer.Option(False, "--version", help="Show version information"),
23
+ ):
24
+ """Janito - AI-powered code modification assistant"""
25
+ if version:
26
+ console = Console()
27
+ console.print(f"Janito version {get_version()}")
28
+ return
29
+
30
+ workdir = workdir or Path.cwd()
31
+ config.set_debug(debug)
32
+ config.set_verbose(verbose)
33
+
34
+ agent = AgentSingleton.get_agent()
35
+
36
+ if ask:
37
+ handle_ask(ask, workdir, include, False, agent)
38
+ elif play:
39
+ handle_play(play, workdir, False)
40
+ elif change_request == "scan":
41
+ paths_to_scan = include if include else [workdir]
42
+ handle_scan(paths_to_scan, workdir)
43
+ elif change_request:
44
+ handle_request(change_request, workdir, include, False, agent)
45
+ else:
46
+ console = Console()
47
+ console.print("Error: Please provide a change request or use --ask/--play options")
48
+ raise typer.Exit(1)
49
+
50
+ def main():
51
+ typer.run(typer_main)
52
+
53
+ if __name__ == "__main__":
54
+ main()
@@ -0,0 +1,113 @@
1
+ from typing import List, Tuple, Optional, NamedTuple
2
+ from difflib import SequenceMatcher
3
+ from janito.config import config
4
+ from rich.console import Console
5
+
6
+ class ContextError(NamedTuple):
7
+ """Contains error details for context matching failures"""
8
+ pre_context: List[str]
9
+ post_context: List[str]
10
+ content: str
11
+
12
+ def parse_change_block(content: str) -> Tuple[List[str], List[str], List[str]]:
13
+ """Parse a change block into pre-context, post-context and change lines.
14
+ Returns (pre_context_lines, post_context_lines, change_lines)"""
15
+ pre_context_lines = []
16
+ post_context_lines = []
17
+ change_lines = []
18
+ in_pre_context = True
19
+
20
+ for line in content.splitlines():
21
+ if line.startswith('='):
22
+ if in_pre_context:
23
+ pre_context_lines.append(line[1:])
24
+ else:
25
+ post_context_lines.append(line[1:])
26
+ elif line.startswith('>'):
27
+ in_pre_context = False
28
+ change_lines.append(line[1:])
29
+
30
+ return pre_context_lines, post_context_lines, change_lines
31
+
32
+ def find_context_match(file_content: str, pre_context: List[str], post_context: List[str], min_context: int = 2) -> Optional[Tuple[int, int]]:
33
+ """Find exact matching location using line-by-line matching.
34
+ Returns (start_index, end_index) or None if no match found."""
35
+ if not (pre_context or post_context) or (len(pre_context) + len(post_context)) < min_context:
36
+ return None
37
+
38
+ file_lines = file_content.splitlines()
39
+
40
+ # Function to check if lines match at a given position
41
+ def lines_match_at(pos: int, target_lines: List[str]) -> bool:
42
+ if pos + len(target_lines) > len(file_lines):
43
+ return False
44
+ return all(a == b for a, b in zip(file_lines[pos:pos + len(target_lines)], target_lines))
45
+
46
+ # For debug output
47
+ debug_matches = []
48
+
49
+ # Try to find pre_context match
50
+ pre_match_pos = None
51
+ if pre_context:
52
+ for i in range(len(file_lines) - len(pre_context) + 1):
53
+ if lines_match_at(i, pre_context):
54
+ pre_match_pos = i
55
+ break
56
+ if config.debug:
57
+ # Record first 20 non-matches for debug output
58
+ if len(debug_matches) < 20:
59
+ debug_matches.append((i, file_lines[i:i + len(pre_context)]))
60
+
61
+ # Try to find post_context match after pre_context if found
62
+ if pre_match_pos is not None and post_context:
63
+ expected_post_pos = pre_match_pos + len(pre_context)
64
+ if not lines_match_at(expected_post_pos, post_context):
65
+ pre_match_pos = None
66
+
67
+ if pre_match_pos is None and config.debug:
68
+ console = Console()
69
+ console.print("\n[bold red]Context Match Debug:[/bold red]")
70
+
71
+ if pre_context:
72
+ console.print("\n[yellow]Expected pre-context:[/yellow]")
73
+ for i, line in enumerate(pre_context):
74
+ console.print(f" {i+1:2d} | '{line}'")
75
+
76
+ if post_context:
77
+ console.print("\n[yellow]Expected post-context:[/yellow]")
78
+ for i, line in enumerate(post_context):
79
+ console.print(f" {i+1:2d} | '{line}'")
80
+
81
+ console.print("\n[yellow]First 20 attempted matches in file:[/yellow]")
82
+ for pos, lines in debug_matches:
83
+ console.print(f"\n[cyan]At line {pos+1}:[/cyan]")
84
+ for i, line in enumerate(lines):
85
+ match_status = "≠" if i < len(pre_context) and line != pre_context[i] else "="
86
+ console.print(f" {i+1:2d} | '{line}' {match_status}")
87
+
88
+ return None
89
+
90
+ if pre_match_pos is None:
91
+ return None
92
+
93
+ end_pos = pre_match_pos + len(pre_context)
94
+
95
+ return pre_match_pos, end_pos
96
+
97
+ def apply_changes(content: str,
98
+ pre_context_lines: List[str],
99
+ post_context_lines: List[str],
100
+ change_lines: List[str]) -> Optional[Tuple[str, Optional[ContextError]]]:
101
+ """Apply changes with context matching, returns (new_content, error_details)"""
102
+ if not content.strip() and not pre_context_lines and not post_context_lines:
103
+ return '\n'.join(change_lines), None
104
+
105
+ pre_context = '\n'.join(pre_context_lines)
106
+ post_context = '\n'.join(post_context_lines)
107
+
108
+ if pre_context and pre_context not in content:
109
+ return None, ContextError(pre_context_lines, post_context_lines, content)
110
+
111
+ if post_context and post_context not in content:
112
+ return None, ContextError(pre_context_lines, post_context_lines, content)
113
+
@@ -0,0 +1,22 @@
1
+ import os
2
+
3
+ SYSTEM_PROMPT = """I am Janito, your friendly software development buddy. I help you with coding tasks while being clear and concise in my responses."""
4
+
5
+ ai_backend = os.getenv('AI_BACKEND', 'claudeai').lower()
6
+
7
+ if ai_backend == 'openai':
8
+ from .openai import OpenAIAgent as AIAgent
9
+ elif ai_backend == 'claudeai':
10
+ from .claudeai import ClaudeAIAgent as AIAgent
11
+ else:
12
+ raise ValueError(f"Unsupported AI_BACKEND: {ai_backend}")
13
+
14
+ class AgentSingleton:
15
+ _instance = None
16
+
17
+ @classmethod
18
+ def get_agent(cls):
19
+ if cls._instance is None:
20
+ cls._instance = AIAgent(SYSTEM_PROMPT)
21
+ return cls._instance
22
+
@@ -0,0 +1,21 @@
1
+
2
+ from abc import ABC, abstractmethod
3
+ from threading import Event
4
+ from typing import Optional, List, Tuple
5
+
6
+ class Agent(ABC):
7
+ """Abstract base class for AI agents"""
8
+ def __init__(self, api_key: Optional[str] = None, system_prompt: str = None):
9
+ self.api_key = api_key
10
+ self.system_message = system_prompt
11
+ self.last_prompt = None
12
+ self.last_full_message = None
13
+ self.last_response = None
14
+ self.messages_history: List[Tuple[str, str]] = []
15
+ if system_prompt:
16
+ self.messages_history.append(("system", system_prompt))
17
+
18
+ @abstractmethod
19
+ def send_message(self, message: str, stop_event: Event = None) -> str:
20
+ """Send message to AI service and return response"""
21
+ pass
@@ -0,0 +1,64 @@
1
+ import anthropic
2
+ import os
3
+ from typing import Optional
4
+ from threading import Event
5
+ from .agent import Agent
6
+
7
+ class ClaudeAIAgent(Agent):
8
+ """Handles interaction with Claude API, including message handling"""
9
+ DEFAULT_MODEL = "claude-3-5-sonnet-20241022"
10
+
11
+ def __init__(self, system_prompt: str = None):
12
+ self.api_key = os.getenv('ANTHROPIC_API_KEY')
13
+ super().__init__(self.api_key, system_prompt)
14
+ if not system_prompt:
15
+ raise ValueError("system_prompt is required")
16
+
17
+ if not self.api_key:
18
+ raise ValueError("ANTHROPIC_API_KEY environment variable is required")
19
+ self.client = anthropic.Client(api_key=self.api_key)
20
+ self.model = os.getenv('CLAUDE_MODEL', self.DEFAULT_MODEL)
21
+ self.system_message = system_prompt
22
+ self.last_prompt = None
23
+ self.last_full_message = None
24
+ self.last_response = None
25
+ self.messages_history = []
26
+ if system_prompt:
27
+ self.messages_history.append(("system", system_prompt))
28
+
29
+ def send_message(self, message: str, stop_event: Event = None) -> str:
30
+ """Send message to Claude API and return response"""
31
+ self.messages_history.append(("user", message))
32
+ # Store the full message
33
+ self.last_full_message = message
34
+
35
+ try:
36
+ # Check if already cancelled
37
+ if stop_event and stop_event.is_set():
38
+ return ""
39
+
40
+ response = self.client.messages.create(
41
+ model=self.model, # Use discovered model
42
+ system=self.system_message,
43
+ max_tokens=8192,
44
+ messages=[
45
+ {"role": "user", "content": message}
46
+ ],
47
+ temperature=0,
48
+ )
49
+
50
+ # Handle response
51
+ response_text = response.content[0].text
52
+
53
+ # Only store and process response if not cancelled
54
+ if not (stop_event and stop_event.is_set()):
55
+ self.last_response = response_text
56
+ self.messages_history.append(("assistant", response_text))
57
+
58
+ # Always return the response, let caller handle cancellation
59
+ return response_text
60
+
61
+ except KeyboardInterrupt:
62
+ if stop_event:
63
+ stop_event.set()
64
+ return ""
@@ -0,0 +1,53 @@
1
+ import openai # updated import
2
+ import os
3
+ from typing import Optional
4
+ from threading import Event
5
+ from .agent import Agent
6
+
7
+ class OpenAIAgent(Agent):
8
+ """Handles interaction with OpenAI API, including message handling"""
9
+ DEFAULT_MODEL = "o1-mini-2024-09-12"
10
+
11
+ def __init__(self, api_key: Optional[str] = None, system_prompt: str = None):
12
+ super().__init__(api_key, system_prompt)
13
+ if not system_prompt:
14
+ raise ValueError("system_prompt is required")
15
+ self.api_key = api_key or os.getenv('OPENAI_API_KEY')
16
+ if not self.api_key:
17
+ raise ValueError("OPENAI_API_KEY environment variable is required")
18
+ openai.api_key = self.api_key
19
+ openai.organization = os.getenv("OPENAI_ORG")
20
+ self.client = openai.Client() # initialized client
21
+ self.model = os.getenv('OPENAI_MODEL', "o1-mini-2024-09-12") # reverted to original default model
22
+
23
+ def send_message(self, message: str, stop_event: Event = None) -> str:
24
+ """Send message to OpenAI API and return response"""
25
+ self.messages_history.append(("user", message))
26
+ self.last_full_message = message
27
+
28
+ try:
29
+ if stop_event and stop_event.is_set():
30
+ return ""
31
+
32
+ #messages = [{"role": "system", "content": self.system_message}]
33
+ messages = [{"role": "user", "content": message}]
34
+
35
+ response = self.client.chat.completions.create(
36
+ model=self.model,
37
+ messages=messages,
38
+ max_completion_tokens=4000,
39
+ temperature=1,
40
+ )
41
+
42
+ response_text = response.choices[0].message.content
43
+
44
+ if not (stop_event and stop_event.is_set()):
45
+ self.last_response = response_text
46
+ self.messages_history.append(("assistant", response_text))
47
+
48
+ return response_text
49
+
50
+ except KeyboardInterrupt:
51
+ if stop_event:
52
+ stop_event.set()
53
+ return ""
@@ -0,0 +1,34 @@
1
+ import unittest
2
+ import os
3
+ from unittest.mock import patch, MagicMock
4
+ from .openai import OpenAIAgent
5
+ from .claudeai import AIAgent
6
+
7
+ class TestAIAgents(unittest.TestCase):
8
+ def setUp(self):
9
+ self.system_prompt = "You are a helpful assistant."
10
+ self.test_message = "Hello, how are you?"
11
+
12
+ def test_openai_agent_initialization(self):
13
+ with patch.dict(os.environ, {'OPENAI_API_KEY': 'test_key'}):
14
+ agent = OpenAIAgent(system_prompt=self.system_prompt)
15
+
16
+ def test_claudeai_agent_initialization(self):
17
+ with patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test_key'}):
18
+ agent = AIAgent(system_prompt=self.system_prompt)
19
+
20
+ def test_openai_agent_send_message(self):
21
+ with patch('openai.OpenAI.chat.completions.create') as mock_create:
22
+ mock_response = MagicMock()
23
+ mock_response.choices[0].message.content = "I'm good, thank you!"
24
+ mock_create.return_value = mock_response
25
+ response = self.openai_agent.send_message(self.test_message)
26
+ self.assertEqual(response, "I'm good, thank you!")
27
+
28
+ def test_claudeai_agent_send_message(self):
29
+ with patch('anthropic.Client.messages.create') as mock_create:
30
+ mock_response = MagicMock()
31
+ mock_response.content[0].text = "I'm Claude, how can I assist you?"
32
+ mock_create.return_value = mock_response
33
+ response = self.claudeai_agent.send_message(self.test_message)
34
+ self.assertEqual(response, "I'm Claude, how can I assist you?")