progress ai forgot to push
This commit is contained in:
@@ -12,6 +12,7 @@ if thirdparty not in sys.path:
|
||||
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
|
||||
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
os.environ["AI_SERVER_ENABLED"] = "1"
|
||||
|
||||
from defer.sugar import install as _install_defer
|
||||
_install_defer()
|
||||
|
||||
+33
-15
@@ -49,6 +49,19 @@ _history_trunc_limit: int = 8000
|
||||
# Global event emitter for API lifecycle events
|
||||
events: EventEmitter = EventEmitter()
|
||||
|
||||
_ai_proxy = None
|
||||
|
||||
def _get_proxy():
|
||||
global _ai_proxy
|
||||
if _ai_proxy is None and os.environ.get("AI_SERVER_ENABLED"):
|
||||
try:
|
||||
from src.ai_client_proxy import AIProxyClient
|
||||
_ai_proxy = AIProxyClient()
|
||||
_ai_proxy.start_server()
|
||||
except Exception:
|
||||
_ai_proxy = None
|
||||
return _ai_proxy
|
||||
|
||||
def set_model_params(temp: float, max_tok: int, trunc_limit: int = 8000, top_p: float = 1.0) -> None:
|
||||
"""
|
||||
Sets global generation parameters like temperature and max tokens.
|
||||
@@ -532,21 +545,26 @@ def get_gemini_cache_stats() -> dict[str, Any]:
|
||||
}
|
||||
|
||||
def list_models(provider: str) -> list[str]:
|
||||
"""
|
||||
[C: src/app_controller.py:AppController.do_fetch, tests/test_agent_capabilities.py:test_agent_capabilities_listing, tests/test_ai_client_list_models.py:test_list_models_gemini_cli, tests/test_deepseek_infra.py:test_deepseek_model_listing, tests/test_minimax_provider.py:test_minimax_list_models]
|
||||
"""
|
||||
creds = _load_credentials()
|
||||
if provider == "gemini":
|
||||
return _list_gemini_models(creds["gemini"]["api_key"])
|
||||
elif provider == "anthropic":
|
||||
return _list_anthropic_models()
|
||||
elif provider == "deepseek":
|
||||
return _list_deepseek_models(creds["deepseek"]["api_key"])
|
||||
elif provider == "gemini_cli":
|
||||
return _list_gemini_cli_models()
|
||||
elif provider == "minimax":
|
||||
return _list_minimax_models(creds["minimax"]["api_key"])
|
||||
return []
|
||||
"""
|
||||
[C: src/app_controller.py:AppController.do_fetch, tests/test_agent_capabilities.py:test_agent_capabilities_listing, tests/test_ai_client_list_models.py:test_list_models_gemini_cli, tests/test_deepseek_infra.py:test_deepseek_model_listing, tests/test_minimax_provider.py:test_minimax_list_models]
|
||||
"""
|
||||
proxy = _get_proxy()
|
||||
if proxy and proxy.status == "ready":
|
||||
result = proxy.send_command("list_models", {"provider": provider})
|
||||
if "result" in result:
|
||||
return result["result"].get("models", [])
|
||||
creds = _load_credentials()
|
||||
if provider == "gemini":
|
||||
return _list_gemini_models(creds["gemini"]["api_key"])
|
||||
elif provider == "anthropic":
|
||||
return _list_anthropic_models()
|
||||
elif provider == "deepseek":
|
||||
return _list_deepseek_models(creds["deepseek"]["api_key"])
|
||||
elif provider == "gemini_cli":
|
||||
return _list_gemini_cli_models()
|
||||
elif provider == "minimax":
|
||||
return _list_minimax_models(creds["minimax"]["api_key"])
|
||||
return []
|
||||
|
||||
def _list_gemini_cli_models() -> list[str]:
|
||||
return [
|
||||
|
||||
@@ -0,0 +1,270 @@
|
||||
import sys
|
||||
import threading
|
||||
import json
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
sys.path.insert(0, '.')
|
||||
|
||||
from src.ai_client_proxy import AIProxyClient
|
||||
|
||||
|
||||
def test_initial_status_is_disconnected():
|
||||
client = AIProxyClient()
|
||||
assert client.status == 'disconnected', f'Expected disconnected, got {client.status}'
|
||||
print('PASS: test_initial_status_is_disconnected')
|
||||
|
||||
|
||||
def test_initial_state_variables():
|
||||
client = AIProxyClient()
|
||||
assert client._process is None, 'Expected _process to be None'
|
||||
assert client._status == 'disconnected', f'Expected disconnected, got {client._status}'
|
||||
assert client._pending == {}, f'Expected empty pending, got {client._pending}'
|
||||
assert client._reader_thread is None, 'Expected _reader_thread to be None'
|
||||
assert hasattr(client, '_pending_lock'), 'Missing _pending_lock'
|
||||
assert callable(getattr(client._pending_lock, 'acquire', None)), 'Expected lock with acquire method'
|
||||
print('PASS: test_initial_state_variables')
|
||||
|
||||
|
||||
def test_status_property():
|
||||
client = AIProxyClient()
|
||||
client._status = 'init'
|
||||
assert client.status == 'init'
|
||||
client._status = 'ready'
|
||||
assert client.status == 'ready'
|
||||
client._status = 'busy'
|
||||
assert client.status == 'busy'
|
||||
client._status = 'error'
|
||||
assert client.status == 'error'
|
||||
client._status = 'disconnected'
|
||||
assert client.status == 'disconnected'
|
||||
print('PASS: test_status_property')
|
||||
|
||||
|
||||
def test_start_server_spawns_subprocess():
|
||||
client = AIProxyClient()
|
||||
with patch('subprocess.Popen') as mock_popen:
|
||||
client.start_server()
|
||||
mock_popen.assert_called_once()
|
||||
args = mock_popen.call_args[0][0]
|
||||
assert '-m' in args, f'Expected -m in args, got {args}'
|
||||
assert 'src.ai_server' in args, f'Expected src.ai_server in args, got {args}'
|
||||
print('PASS: test_start_server_spawns_subprocess')
|
||||
|
||||
|
||||
def test_start_server_sets_status_to_init():
|
||||
client = AIProxyClient()
|
||||
with patch('subprocess.Popen'):
|
||||
client.start_server()
|
||||
assert client._status == 'init', f'Expected init, got {client._status}'
|
||||
print('PASS: test_start_server_sets_status_to_init')
|
||||
|
||||
|
||||
def test_start_server_starts_reader_thread():
|
||||
client = AIProxyClient()
|
||||
with patch('subprocess.Popen'):
|
||||
client.start_server()
|
||||
assert client._reader_thread is not None, 'Expected reader thread to be started'
|
||||
assert isinstance(client._reader_thread, threading.Thread), 'Expected threading.Thread'
|
||||
print('PASS: test_start_server_starts_reader_thread')
|
||||
|
||||
|
||||
def test_send_command_includes_method_and_params():
|
||||
client = AIProxyClient()
|
||||
client._process = MagicMock()
|
||||
client._status = 'ready'
|
||||
client._pending_lock = threading.Lock()
|
||||
client._reader_thread = MagicMock()
|
||||
written_data = {}
|
||||
def capture_write(data):
|
||||
nonlocal written_data
|
||||
written_data = json.loads(data)
|
||||
client._write_to_stdin = capture_write
|
||||
client._pending = {}
|
||||
client.send_command('test_method', {'param': 'value'})
|
||||
assert written_data['id'], f"Expected id to be set"
|
||||
assert written_data['method'] == 'test_method', f"Expected method test_method, got {written_data.get('method')}"
|
||||
assert written_data['params'] == {'param': 'value'}, f"Expected params, got {written_data.get('params')}"
|
||||
print('PASS: test_send_command_includes_method_and_params')
|
||||
|
||||
|
||||
def test_send_command_adds_pending_request():
|
||||
client = AIProxyClient()
|
||||
client._process = MagicMock()
|
||||
client._status = 'ready'
|
||||
client._pending_lock = threading.Lock()
|
||||
client._reader_thread = MagicMock()
|
||||
client._pending = {}
|
||||
client._write_to_stdin = lambda x: None
|
||||
client.send_command('test_method', {})
|
||||
assert len(client._pending) == 1, f"Expected 1 pending request, got {len(client._pending)}"
|
||||
req_id = list(client._pending.keys())[0]
|
||||
assert 'event' in client._pending[req_id], f"Expected 'event' key in pending request"
|
||||
print('PASS: test_send_command_adds_pending_request')
|
||||
|
||||
|
||||
def test_send_command_waits_for_event():
|
||||
client = AIProxyClient()
|
||||
client._process = MagicMock()
|
||||
client._status = 'ready'
|
||||
client._pending_lock = threading.Lock()
|
||||
client._reader_thread = MagicMock()
|
||||
client._pending = {}
|
||||
wait_called = [False]
|
||||
def mock_wait(timeout):
|
||||
wait_called[0] = True
|
||||
return True
|
||||
mock_event = MagicMock()
|
||||
mock_event.wait = mock_wait
|
||||
client._write_to_stdin = lambda x: None
|
||||
original_event = threading.Event
|
||||
with patch('threading.Event', return_value=mock_event):
|
||||
client.send_command('test_method', {})
|
||||
assert wait_called[0], "Expected wait to be called"
|
||||
print('PASS: test_send_command_waits_for_event')
|
||||
|
||||
|
||||
def test_send_command_returns_response_when_ready():
|
||||
client = AIProxyClient()
|
||||
client._process = MagicMock()
|
||||
client._status = 'ready'
|
||||
client._pending_lock = threading.Lock()
|
||||
client._reader_thread = MagicMock()
|
||||
test_response = {'id': 'test-uuid-123', 'result': {'output': 'success'}}
|
||||
mock_event = MagicMock()
|
||||
mock_event.wait = lambda timeout: True
|
||||
client._pending = {'test-uuid-123': {'event': mock_event, 'response': test_response}}
|
||||
client._write_to_stdin = lambda x: None
|
||||
result = client.send_command('test_method', {})
|
||||
assert result == test_response, f'Expected {test_response}, got {result}'
|
||||
print('PASS: test_send_command_returns_response_when_ready')
|
||||
|
||||
|
||||
def test_send_command_timeout_raises():
|
||||
client = AIProxyClient()
|
||||
client._process = MagicMock()
|
||||
client._status = 'ready'
|
||||
client._pending_lock = threading.Lock()
|
||||
client._reader_thread = MagicMock()
|
||||
mock_event = MagicMock()
|
||||
mock_event.wait = lambda timeout: False
|
||||
client._pending = {'test-uuid-123': {'event': mock_event, 'response': None}}
|
||||
client._write_to_stdin = lambda x: None
|
||||
try:
|
||||
client.send_command('test_method', {})
|
||||
assert False, 'Expected TimeoutError'
|
||||
except TimeoutError as e:
|
||||
assert 'timed out' in str(e), f'Expected timed out message, got {e}'
|
||||
print('PASS: test_send_command_timeout_raises')
|
||||
|
||||
|
||||
def test_send_command_cleans_pending_on_timeout():
|
||||
client = AIProxyClient()
|
||||
client._process = MagicMock()
|
||||
client._status = 'ready'
|
||||
client._pending_lock = threading.Lock()
|
||||
client._reader_thread = MagicMock()
|
||||
mock_event = MagicMock()
|
||||
mock_event.wait = lambda timeout: False
|
||||
client._pending = {'test-uuid-123': {'event': mock_event, 'response': None}}
|
||||
client._write_to_stdin = lambda x: None
|
||||
try:
|
||||
client.send_command('test_method', {})
|
||||
except TimeoutError:
|
||||
pass
|
||||
assert 'test-uuid-123' not in client._pending, f'Expected cleanup on timeout, got {client._pending}'
|
||||
print('PASS: test_send_command_cleans_pending_on_timeout')
|
||||
|
||||
|
||||
def test_read_loop_parses_json_and_sets_response():
|
||||
client = AIProxyClient()
|
||||
client._pending = {'test-uuid': {'event': MagicMock(), 'response': None}}
|
||||
client._pending_lock = threading.Lock()
|
||||
client._process = MagicMock()
|
||||
lines = [b'{"id": "test-uuid", "result": {"data": "test"}}']
|
||||
client._process.stdout = MagicMock()
|
||||
client._process.stdout.readline.side_effect = lines + [b'']
|
||||
client._read_loop()
|
||||
assert client._pending['test-uuid']['response'] == {'id': 'test-uuid', 'result': {'data': 'test'}}, f"Expected parsed response"
|
||||
client._pending['test-uuid']['event'].set.assert_called_once()
|
||||
print('PASS: test_read_loop_parses_json_and_sets_response')
|
||||
|
||||
|
||||
def test_read_loop_matches_by_id():
|
||||
client = AIProxyClient()
|
||||
client._pending = {
|
||||
'uuid-1': {'event': MagicMock(), 'response': None},
|
||||
'uuid-2': {'event': MagicMock(), 'response': None}
|
||||
}
|
||||
client._pending_lock = threading.Lock()
|
||||
client._process = MagicMock()
|
||||
lines = [b'{"id": "uuid-2", "result": {"data": "from-uuid2"}}']
|
||||
client._process.stdout = MagicMock()
|
||||
client._process.stdout.readline.side_effect = lines + [b'']
|
||||
client._read_loop()
|
||||
client._pending['uuid-1']['event'].set.assert_not_called()
|
||||
client._pending['uuid-2']['event'].set.assert_called_once()
|
||||
print('PASS: test_read_loop_matches_by_id')
|
||||
|
||||
|
||||
def test_stop_terminates_process():
|
||||
client = AIProxyClient()
|
||||
mock_process = MagicMock()
|
||||
client._process = mock_process
|
||||
client._status = 'ready'
|
||||
client._pending = {'test': {'event': MagicMock(), 'response': {}}}
|
||||
client._pending_lock = threading.Lock()
|
||||
client.stop()
|
||||
mock_process.terminate.assert_called_once()
|
||||
print('PASS: test_stop_terminates_process')
|
||||
|
||||
|
||||
def test_stop_cleans_pending():
|
||||
client = AIProxyClient()
|
||||
mock_process = MagicMock()
|
||||
client._process = mock_process
|
||||
client._status = 'ready'
|
||||
client._pending = {'test': {'event': MagicMock(), 'response': {}}}
|
||||
client._pending_lock = threading.Lock()
|
||||
client.stop()
|
||||
assert client._pending == {}, f'Expected empty pending after stop, got {client._pending}'
|
||||
print('PASS: test_stop_cleans_pending')
|
||||
|
||||
|
||||
def test_stop_sets_status_to_disconnected():
|
||||
client = AIProxyClient()
|
||||
mock_process = MagicMock()
|
||||
client._process = mock_process
|
||||
client._status = 'ready'
|
||||
client._pending = {'test': {'event': MagicMock(), 'response': {}}}
|
||||
client._pending_lock = threading.Lock()
|
||||
client.stop()
|
||||
assert client._status == 'disconnected', f'Expected disconnected, got {client._status}'
|
||||
print('PASS: test_stop_sets_status_to_disconnected')
|
||||
|
||||
|
||||
def test_default_timeout_is_60_seconds():
|
||||
client = AIProxyClient()
|
||||
assert client._DEFAULT_TIMEOUT == 60.0, f'Expected 60.0, got {client._DEFAULT_TIMEOUT}'
|
||||
print('PASS: test_default_timeout_is_60_seconds')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_initial_status_is_disconnected()
|
||||
test_initial_state_variables()
|
||||
test_status_property()
|
||||
test_start_server_spawns_subprocess()
|
||||
test_start_server_sets_status_to_init()
|
||||
test_start_server_starts_reader_thread()
|
||||
test_send_command_includes_method_and_params()
|
||||
test_send_command_adds_pending_request()
|
||||
test_send_command_waits_for_event()
|
||||
test_send_command_returns_response_when_ready()
|
||||
test_send_command_timeout_raises()
|
||||
test_send_command_cleans_pending_on_timeout()
|
||||
test_read_loop_parses_json_and_sets_response()
|
||||
test_read_loop_matches_by_id()
|
||||
test_stop_terminates_process()
|
||||
test_stop_cleans_pending()
|
||||
test_stop_sets_status_to_disconnected()
|
||||
test_default_timeout_is_60_seconds()
|
||||
print('\nAll tests passed!')
|
||||
Reference in New Issue
Block a user