Atlas / tests /local_app.py
findEthics
Adding analytics capabilities
1d4dc07
Raw
History Blame Contribute Delete
4.06 kB
#!/usr/bin/env python3
"""
Local interactive chat client for Atlas API
Connects to the FastAPI server running on localhost:7860
"""
import httpx
import asyncio
import sys
import json
API_BASE_URL = "http://localhost:7860"
# Global session tracking
_session_id = None
async def send_chat_message(prompt: str, use_search: bool = True) -> dict:
"""Send a chat message to the API and return the response"""
global _session_id
async with httpx.AsyncClient(timeout=30.0) as client:
try:
# Prepare headers with session ID if available
headers = {
"User-Agent": "local_app.py/1.0",
"Content-Type": "application/json"
}
if _session_id:
headers["X-Session-ID"] = _session_id
response = await client.post(
f"{API_BASE_URL}/chat",
json={
"prompt": prompt,
"max_new_tokens": 500,
"use_search": use_search,
"temperature": 0.7
},
headers=headers
)
response.raise_for_status()
# Extract session ID from response headers if available
if "X-Session-ID" in response.headers:
_session_id = response.headers["X-Session-ID"]
return response.json()
except httpx.ConnectError:
return {"error": "Could not connect to API server. Make sure app.py is running on localhost:7860"}
except httpx.TimeoutException:
return {"error": "Request timed out. Please try again."}
except httpx.HTTPStatusError as e:
return {"error": f"API error: {e.response.status_code} - {e.response.text}"}
except Exception as e:
return {"error": f"Unexpected error: {str(e)}"}
async def main():
"""Main chat loop"""
global _session_id
print("Atlas Chat Client with Analytics")
print("Type your messages and press Enter. Empty input or Ctrl+C to exit.")
print("Type 'search:' before your message to enable web search.")
print("-" * 50)
try:
while True:
try:
# Get user input
user_input = input("\nYou: ").strip()
# Exit on empty input
if not user_input:
print("Goodbye!")
break
# Check for search prefix
use_search = True
if user_input.startswith("search:"):
use_search = True
user_input = user_input[7:].strip()
if not user_input:
print("Please provide a message after 'search:'")
continue
# Send message to API
print("AI: ", end="", flush=True)
response = await send_chat_message(user_input, use_search)
# Handle response
if "error" in response:
print(f"Error: {response['error']}")
else:
print(response.get("response", "No response received"))
# Show search results if available
if response.get("search_results"):
print(f"\n[Found {len(response['search_results'])} search results]")
# Show session info on first message
if _session_id:
print(f"\n[Session: {_session_id[:8]}...]")
except KeyboardInterrupt:
print("\nGoodbye!")
break
except EOFError:
print("\nGoodbye!")
break
except Exception as e:
print(f"Fatal error: {e}")
sys.exit(1)
if __name__ == "__main__":
asyncio.run(main())