ai/service.py

import json
import openai
import requests
from pydantic import ValidationError
from models.market_data import MarketData
from models.suggestion import (
    TradeSuggestion,
    TradeDirection,
    TakeProfitPoints,
    RecommendationType,
)


class AIService:
    @staticmethod
    def _prepare_prompt(
        symbol, leverage, trade_amount, market_data: MarketData, current_price: float
    ):
        klines = [
            {
                "timestamp": k.timestamp,
                "open": k.open,
                "high": k.high,
                "low": k.low,
                "close": k.close,
                "volume": k.volume,
            }
            for k in market_data.klines
        ]

        return f"""
You are an expert crypto futures trader. Based on market data, suggest a trade for {symbol}.

Current Price: ${current_price}
Leverage: {leverage}x
Amount: ${trade_amount}

Market Data:
{json.dumps(klines, indent=2)}

Only return valid JSON:
{{
  "direction": "long" or "short",
  "entry_price": float,
  "recommended_leverage": int (10-75),
  "take_profit": {{
    "first": float,
    "second": float,
    "third": float
  }},
  "recommendation": "It is recommended to enter the transaction." or "It is not recommended to enter into a transaction." or "Entering the trade with caution"
}}
"""

    def _parse_response(
        self, response: str, symbol: str, current_price: float, trade_amount: float
    ):
        try:
            if response.startswith("```json"):
                response = response[7:]
            if response.endswith("```"):
                response = response[:-3]
            data = json.loads(response)

            rec_map = {
                "It is recommended to enter the transaction.": RecommendationType.RECOMMENDED,
                "It is not recommended to enter into a transaction.": RecommendationType.NOT_RECOMMENDED,
                "Entering the trade with caution": RecommendationType.CAUTIOUS,
            }

            return TradeSuggestion(
                symbol=symbol,
                direction=TradeDirection(data["direction"]),
                entry_price=float(data["entry_price"]),
                recommended_leverage=int(data["recommended_leverage"]),
                take_profit=TakeProfitPoints(**data["take_profit"]),
                recommendation=rec_map.get(
                    data["recommendation"], RecommendationType.CAUTIOUS
                ),
                current_price=current_price,
                trade_amount=trade_amount,
            )
        except Exception as e:
            raise ValueError(f"Failed to parse AI response: {e}\nRaw: {response}")

    def generate(
        self,
        symbol,
        leverage,
        trade_amount,
        market_data,
        current_price,
        provider,
        openai_key,
        hf_token,
    ):
        prompt = self._prepare_prompt(
            symbol, leverage, trade_amount, market_data, current_price
        )

        if provider == "OpenAI":
            openai.api_key = openai_key
            try:
                res = openai.ChatCompletion.create(
                    model="gpt-3.5-turbo",
                    messages=[
                        {
                            "role": "system",
                            "content": "You are a crypto trading expert.",
                        },
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.2,
                )
                content = res.choices[0].message.content
                return self._parse_response(
                    content, symbol, current_price, trade_amount
                )
            except Exception as e:
                print(f"OpenAI Error: {e}")
                return None

        elif provider == "HuggingFace":
            try:
                headers = {"Authorization": f"Bearer {hf_token}"}
                data = {"inputs": prompt}
                url = "https://api-inference.huggingface.co/models/tiiuae/falcon-7b"
                res = requests.post(url, headers=headers, json=data)
                response = res.json()
                text = response[0]["generated_text"]
                return self._parse_response(text, symbol, current_price, trade_amount)
            except Exception as e:
                print(f"HuggingFace Error: {e}")
                return None