api_wrapper.py

#!/usr/bin/env python3
import http.server
import json
import os
import sys
import logging
from pathlib import Path
import time
from urllib.parse import parse_qs, urlparse

# Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger("ten_api")

# Проверяем, запущены ли мы в HuggingFace Space
IS_HF_SPACE = os.environ.get("SPACE_ID") is not None
USE_WRAPPER = os.environ.get("USE_WRAPPER", "false").lower() in ("true", "1", "yes")

# Путь к директории с агентами
if IS_HF_SPACE or USE_WRAPPER:
    TMP_DIR = os.environ.get("TMP_DIR", "/tmp/ten_user")
    AGENT_DIR = os.environ.get("TEN_AGENT_DIR", f"{TMP_DIR}/agents")
else:
    AGENT_DIR = os.environ.get("TEN_AGENT_DIR", "/tmp/ten_user/agents")

logger.info(f"Using agent directory: {AGENT_DIR}")
logger.info(f"Running in HuggingFace Space: {IS_HF_SPACE}")
logger.info(f"Using Wrapper: {USE_WRAPPER}")

class TENAgentHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        """Переопределение логирования для вывода в stdout"""
        logger.info("%s - %s", self.address_string(), format % args)
        
    def _set_headers(self, content_type="application/json"):
        self.send_response(200)
        self.send_header('Content-type', content_type)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization')
        self.end_headers()
        
    def do_OPTIONS(self):
        self._set_headers()
        
    def do_GET(self):
        logger.info(f"GET request: {self.path}")
        
        # Разбор URL для параметров
        parsed_url = urlparse(self.path)
        path = parsed_url.path
        query_params = parse_qs(parsed_url.query)
        
        # Обрабатываем API endpoints
        if path == "/graphs" or path == "/api/graphs":
            self._handle_graphs()
        elif path == "/health" or path == "/":
            self._handle_health()
        elif path == "/list":
            self._handle_list()
        elif path == "/dev-tmp/addons/default-properties":
            self._handle_default_properties()
        elif path == "/vector/document/preset/list":
            self._handle_vector_preset_list()
        elif path.startswith("/api/dev/v1/packages"):
            self._handle_packages()
        elif path.startswith("/api/designer/v1/packages"):
            self._handle_packages()
        else:
            # Для всех остальных запросов возвращаем 404
            self.send_error(404, f"Not found: {path}")
    
    def _handle_graphs(self):
        """Обработка запроса на получение списка графов"""
        try:
            property_file = Path(AGENT_DIR) / "property.json"
            if not property_file.exists():
                logger.error(f"Property file not found at {property_file}")
                self.send_error(404, "Property file not found")
                return
            
            with open(property_file, "r") as f:
                property_data = json.load(f)
            
            graphs = property_data.get("graphs", [])
            self._set_headers()
            self.wfile.write(json.dumps(graphs).encode())
            logger.info(f"Returned {len(graphs)} graphs")
        except Exception as e:
            logger.error(f"Error reading property.json: {e}")
            self.send_error(500, f"Internal error: {e}")
    
    def _handle_health(self):
        """Обработка запроса на проверку статуса сервера"""
        self._set_headers()
        self.wfile.write(json.dumps({
            "status": "ok", 
            "time": time.time(), 
            "is_hf_space": IS_HF_SPACE,
            "using_wrapper": True
        }).encode())
    
    def _handle_list(self):
        """Обработка запроса на получение списка активных сессий"""
        self._set_headers()
        self.wfile.write(json.dumps([]).encode())
    
    def _handle_default_properties(self):
        """Обработка запроса на получение настроек по умолчанию"""
        self._set_headers()
        self.wfile.write(json.dumps({}).encode())
    
    def _handle_vector_preset_list(self):
        """Обработка запроса на получение списка пресетов векторов"""
        self._set_headers()
        self.wfile.write(json.dumps([]).encode())
    
    def _handle_packages(self):
        """Обработка запросов к пакетам"""
        # Этот метод эмулирует возврат списка графов и для других эндпоинтов
        try:
            property_file = Path(AGENT_DIR) / "property.json"
            if not property_file.exists():
                logger.error(f"Property file not found at {property_file}")
                self.send_error(404, "Property file not found")
                return
            
            with open(property_file, "r") as f:
                property_data = json.load(f)
            
            graphs = property_data.get("graphs", [])
            self._set_headers()
            response_data = {
                "data": graphs,
                "status": 200,
                "message": "Success"
            }
            self.wfile.write(json.dumps(response_data).encode())
            logger.info(f"Handled packages request and returned {len(graphs)} graphs")
        except Exception as e:
            logger.error(f"Error handling packages request: {e}")
            self.send_error(500, f"Internal error: {e}")
    
    def do_POST(self):
        logger.info(f"POST request: {self.path}")
        
        # Читаем тело запроса
        content_length = int(self.headers['Content-Length']) if 'Content-Length' in self.headers else 0
        post_data = self.rfile.read(content_length)
        
        try:
            request_data = json.loads(post_data) if content_length > 0 else {}
        except json.JSONDecodeError:
            request_data = {}
            
        logger.info(f"Request data: {json.dumps(request_data)[:200]}...")
        
        # Обрабатываем различные POST запросы
        if self.path == "/ping":
            self._handle_ping()
        elif self.path == "/token/generate":
            self._handle_token_generate(request_data)
        elif self.path == "/start":
            self._handle_start(request_data)
        elif self.path == "/stop":
            self._handle_stop(request_data)
        elif self.path == "/vector/document/update" or self.path == "/vector/document/upload":
            self._handle_vector_document(request_data)
        elif self.path.startswith("/api/dev/v1/packages") or self.path.startswith("/api/designer/v1/packages"):
            self._handle_packages_post(request_data)
        else:
            # Для всех остальных запросов возвращаем 404
            self.send_error(404, f"Not found: {self.path}")
    
    def _handle_ping(self):
        """Обработка ping запроса"""
        self._set_headers()
        self.wfile.write(json.dumps({
            "status": "ok", 
            "timestamp": time.time(),
            "server": "ten-agent-api-wrapper",
            "in_hf_space": IS_HF_SPACE
        }).encode())
    
    def _handle_token_generate(self, request_data):
        """Обработка запроса на генерацию токена"""
        self._set_headers()
        response = {
            "token": "dummy_token_for_agora",
            "request_id": request_data.get("RequestId", ""),
            "channel_name": request_data.get("ChannelName", ""),
            "uid": request_data.get("Uid", 0)
        }
        self.wfile.write(json.dumps(response).encode())
        logger.info(f"Generated token for channel: {request_data.get('ChannelName', '')}")
    
    def _handle_start(self, request_data):
        """Обработка запроса на запуск сессии"""
        self._set_headers()
        # Возвращаем успешный статус и ID сессии
        response = {
            "status": "ok",
            "session_id": f"dummy_session_{int(time.time())}",
            "message": "Session started successfully",
            "graph_file": request_data.get("graph_file", "unknown")
        }
        self.wfile.write(json.dumps(response).encode())
        logger.info(f"Started session with graph: {request_data.get('graph_file', '')}")
    
    def _handle_stop(self, request_data):
        """Обработка запроса на остановку сессии"""
        self._set_headers()
        self.wfile.write(json.dumps({
            "status": "ok",
            "message": "Session stopped successfully"
        }).encode())
        logger.info(f"Stopped session: {request_data.get('session_id', '')}")
    
    def _handle_vector_document(self, request_data):
        """Обработка запроса на работу с векторными документами"""
        self._set_headers()
        self.wfile.write(json.dumps({
            "status": "ok",
            "document_id": f"dummy_doc_{int(time.time())}",
            "message": "Document processed successfully"
        }).encode())
        logger.info(f"Processed vector document: {request_data.get('name', 'unnamed')}")
    
    def _handle_packages_post(self, request_data):
        """Обработка POST запросов к пакетам"""
        self._set_headers()
        response_data = {
            "data": {},
            "status": 200,
            "message": "Success"
        }
        self.wfile.write(json.dumps(response_data).encode())
        logger.info(f"Handled packages POST request")

def run(server_class=http.server.HTTPServer, handler_class=TENAgentHandler, port=8080):
    server_address = ('', port)
    httpd = server_class(server_address, handler_class)
    logger.info(f"Starting API server on port {port}...")
    logger.info(f"Using agent directory: {AGENT_DIR}")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Server error: {e}")
        raise

if __name__ == "__main__":
    port = int(os.environ.get("API_PORT", 8080))
    try:
        run(port=port)
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Server error: {e}")