Update ai/service.py

ai/service.py

@@ -1,7 +1,7 @@
 import json
 import openai
 import requests
-from typing import Optional
+from pydantic import ValidationError
 from models.market_data import MarketData
 from models.suggestion import (
     TradeSuggestion,
@@ -12,113 +12,123 @@ from models.suggestion import (
 
 
 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
+    @staticmethod
+    def _prepare_prompt(
+        symbol, leverage, trade_amount, market_data: MarketData, current_price: float
     ):
-        candles = [vars(c) for c in market_data.klines]
+        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 experienced crypto futures trader. Analyze the following:
-Symbol: {symbol}
+You are an expert crypto futures trader. Based on market data, suggest a trade for {symbol}.
+
 Current Price: ${current_price}
-Leverage: {leverage}
-Trade Amount: ${trade_amount}
-Candles: {json.dumps(candles, indent=2)}
+Leverage: {leverage}x
+Amount: ${trade_amount}
 
-Provide output JSON:
+Market Data:
+{json.dumps(klines, indent=2)}
+
+Only return valid JSON:
 {{
-  "direction": "long OR short",
+  "direction": "long" or "short",
   "entry_price": float,
-  "recommended_leverage": 10-75,
+  "recommended_leverage": int (10-75),
   "take_profit": {{
     "first": float,
     "second": float,
     "third": float
   }},
-  "recommendation": "... one of: Recommended / Not recommended / Caution ..."
+  "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(
-        self, raw: str, symbol: str, current_price: float, trade_amount: float
-    ) -> TradeSuggestion:
+    def _parse_response(
+        self, response: str, symbol: str, current_price: float, trade_amount: float
+    ):
         try:
-            raw = raw.strip()
-            if raw.startswith("```json"):
-                raw = raw[7:]
-            if raw.endswith("```"):
-                raw = raw[:-3]
+            if response.startswith("```json"):
+                response = response[7:]
+            if response.endswith("```"):
+                response = response[:-3]
+            data = json.loads(response)
 
-            data = json.loads(raw)
-            recommendation_map = {
+            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,
             }
 
-            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()),
+                direction=TradeDirection(data["direction"]),
                 entry_price=float(data["entry_price"]),
                 recommended_leverage=int(data["recommended_leverage"]),
-                take_profit=take_profit,
-                recommendation=recommendation_map.get(
+                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:
-            print(f"Parse error: {e}")
-            return None
+            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