Create ai/service.py

ai/service.py

@@ -0,0 +1,124 @@
+import json
+import openai
+import requests
+from typing import Optional
+from models.market_data import MarketData
+from models.suggestion import (
+    TradeSuggestion,
+    TradeDirection,
+    TakeProfitPoints,
+    RecommendationType,
+)
+
+
+class AIService:
+    def generate(
+        self,
+        symbol: str,
+        leverage: int,
+        trade_amount: float,
+        current_price: float,
+        market_data: MarketData,
+        provider: str,
+        openai_key: Optional[str] = None,
+        hf_token: Optional[str] = None,
+    ) -> Optional[TradeSuggestion]:
+        prompt = self._build_prompt(
+            symbol, leverage, trade_amount, current_price, market_data
+        )
+
+        try:
+            if provider == "OpenAI":
+                openai.api_key = openai_key
+                response = openai.ChatCompletion.create(
+                    model="gpt-4",
+                    messages=[
+                        {
+                            "role": "system",
+                            "content": "You are an expert crypto futures trader.",
+                        },
+                        {"role": "user", "content": prompt},
+                    ],
+                    temperature=0.2,
+                    max_tokens=600,
+                )
+                content = response["choices"][0]["message"]["content"]
+            else:
+                headers = {"Authorization": f"Bearer {hf_token}"}
+                body = {"inputs": prompt}
+                response = requests.post(
+                    "https://api-inference.huggingface.co/models/tiiuae/falcon-7b",
+                    headers=headers,
+                    json=body,
+                )
+                content = response.json().get("generated_text", "")
+
+            return self._parse(content, symbol, current_price, trade_amount)
+        except Exception as e:
+            print(f"AI error: {e}")
+            return None
+
+    def _build_prompt(
+        self, symbol, leverage, trade_amount, current_price, market_data: MarketData
+    ):
+        candles = [vars(c) for c in market_data.klines]
+        return f"""
+You are an experienced crypto futures trader. Analyze the following:
+Symbol: {symbol}
+Current Price: ${current_price}
+Leverage: {leverage}
+Trade Amount: ${trade_amount}
+Candles: {json.dumps(candles, indent=2)}
+
+Provide output JSON:
+{{
+  "direction": "long OR short",
+  "entry_price": float,
+  "recommended_leverage": 10-75,
+  "take_profit": {{
+    "first": float,
+    "second": float,
+    "third": float
+  }},
+  "recommendation": "... one of: Recommended / Not recommended / Caution ..."
+}}
+        """
+
+    def _parse(
+        self, raw: str, symbol: str, current_price: float, trade_amount: float
+    ) -> TradeSuggestion:
+        try:
+            raw = raw.strip()
+            if raw.startswith("```json"):
+                raw = raw[7:]
+            if raw.endswith("```"):
+                raw = raw[:-3]
+
+            data = json.loads(raw)
+            recommendation_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,
+            }
+
+            take_profit = TakeProfitPoints(
+                first=float(data["take_profit"]["first"]),
+                second=float(data["take_profit"]["second"]),
+                third=float(data["take_profit"]["third"]),
+            )
+
+            return TradeSuggestion(
+                symbol=symbol,
+                direction=TradeDirection(data["direction"].lower()),
+                entry_price=float(data["entry_price"]),
+                recommended_leverage=int(data["recommended_leverage"]),
+                take_profit=take_profit,
+                recommendation=recommendation_map.get(
+                    data["recommendation"], RecommendationType.CAUTIOUS
+                ),
+                current_price=current_price,
+                trade_amount=trade_amount,
+            )
+        except Exception as e:
+            print(f"Parse error: {e}")
+            return None