Update test.py

test.py

@@ -0,0 +1,102 @@
+import requests
+from langchain.llms.base import LLM
+from typing import Optional, List
+from pydantic import Field
+import json
+
+
+class HumbleBeeLLM(LLM):
+    api_base: str = "https://humblebeeai-llm-host.hf.space"
+    username: str = Field(default=None, exclude=True)  # Exclude from Pydantic validation
+    password: str = Field(default=None, exclude=True)
+    token: Optional[str] = None
+
+    def __init__(self, username: str, password: str, **kwargs):
+        super().__init__(**kwargs)  # Initialize LangChain's LLM class properly
+        self.username = username
+        self.password = password
+        self.authenticate()
+
+    def authenticate(self):
+        """Logs in to get the token."""
+        login_payload = {
+            "username": self.username,
+            "password": self.password
+        }
+
+        try:
+            response = requests.post(
+                f"{self.api_base}/login",
+                data=login_payload,  # Must be 'data' instead of 'json' for form-urlencoded
+                headers={"Content-Type": "application/x-www-form-urlencoded"}
+            )
+
+            if response.status_code == 200:
+                self.token = response.json().get("access_token")
+            else:
+                print(f"Login failed: {response.status_code} - {response.text}")
+                raise Exception("Failed to authenticate.")
+
+        except requests.exceptions.RequestException as e:
+            raise Exception(f"Error during authentication: {str(e)}")
+
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        """Sends a request to the LLM API and handles streaming responses."""
+        if not self.token:
+            self.authenticate()
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.token}"
+        }
+        payload = {"query": prompt}
+
+        try:
+            response = requests.post(f"{self.api_base}/generate?stream=true", json=payload, headers=headers, stream=True)
+
+            if response.status_code == 200:
+                # Read the streaming response and reconstruct it token by token
+                full_response = ""
+                for line in response.iter_lines():
+                    if line:
+                        try:
+                            # Clean up the streamed JSON line
+                            data = json.loads(line.decode("utf-8").replace("data: ", ""))
+                            token_content = data.get("content", "")
+                            full_response += token_content
+                            print(token_content, end="", flush=True)  # Print each token as it arrives
+                        except json.JSONDecodeError:
+                            continue  # Ignore invalid JSON lines
+
+                return full_response.strip()  # Return the full response after all tokens
+
+            else:
+                print(f"Error {response.status_code}: {response.text}")
+                return f"Error {response.status_code}: {response.text}"
+
+        except requests.exceptions.RequestException as e:
+            return f"Request error: {str(e)}"
+
+    @property
+    def _identifying_params(self):
+        return {"api_base": self.api_base}
+
+    @property
+    def _llm_type(self):
+        return "humblebee-llm"
+
+
+# ✅ Usage Example
+llm = HumbleBeeLLM(username="testuser", password="testpassword")
+
+"""# Invoke with LangChain's API
+while True:
+    query = input("\n...\n")
+    
+    if query.lower() == 'bye':  # Case insensitive check
+        print("Goodbye!")
+        break  # Exit loop when 'bye' is entered
+    
+    response = llm.invoke(query)
+
+"""