Spaces:
Sleeping
Sleeping
| import re | |
| from dotenv import load_dotenv | |
| import requests | |
| from uuid import uuid4 | |
| import json | |
| from typing import Any, AsyncGenerator, Dict | |
| import os | |
| # Load environment variables from .env file | |
| load_dotenv() | |
| class v2: | |
| def __init__( | |
| self, | |
| timeout: int = 100, | |
| proxies: dict = {}, | |
| ): | |
| self.session = requests.Session() | |
| self.chat_endpoint = os.getenv("v2") | |
| self.stream_chunk_size = 64 | |
| self.timeout = timeout | |
| self.last_response = {} | |
| self.headers = { | |
| "accept": "*/*", | |
| "accept-encoding": "gzip, deflate, br, zstd", | |
| "accept-language": "en-US,en;q=0.9,en-IN;q=0.8", | |
| "content-type": "application/json", | |
| } | |
| self.session.headers.update(self.headers) | |
| self.session.proxies = proxies | |
| def ask( | |
| self, | |
| prompt: str, | |
| stream: bool = False, | |
| raw: bool = False, | |
| ) -> dict: | |
| conversation_prompt =f"Tell me evrything about {prompt}" | |
| self.session.headers.update(self.headers) | |
| payload = { | |
| "query": conversation_prompt, | |
| "search_uuid": uuid4().hex, | |
| "lang": "", | |
| "agent_lang": "en", | |
| "search_options": { | |
| "langcode": "en-US" | |
| }, | |
| "search_video": True, | |
| "contexts_from": "google" | |
| } | |
| def for_stream(): | |
| response = self.session.post( | |
| self.chat_endpoint, json=payload, stream=True, timeout=self.timeout | |
| ) | |
| if not response.ok: | |
| raise Exception( | |
| f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}" | |
| ) | |
| streaming_text = "" | |
| for line in response.iter_lines(decode_unicode=True): | |
| if line.startswith('data:'): | |
| try: | |
| data = json.loads(line[5:].strip()) | |
| if data['type'] == 'answer' and 'text' in data['data']: | |
| new_text = data['data']['text'] | |
| if len(new_text) > len(streaming_text): | |
| delta = new_text[len(streaming_text):] | |
| streaming_text = new_text | |
| resp = dict(text=delta) | |
| self.last_response.update(dict(text=streaming_text)) | |
| yield line if raw else resp | |
| except json.JSONDecodeError: | |
| pass | |
| def for_non_stream(): | |
| full_response = "" | |
| for chunk in for_stream(): | |
| if not raw: | |
| full_response += chunk['text'] | |
| self.last_response = dict(text=full_response) | |
| return self.last_response | |
| return for_stream() if stream else for_non_stream() | |
| def chat( | |
| self, | |
| prompt: str, | |
| stream: bool = False, | |
| ) -> str: | |
| def for_stream(): | |
| for response in self.ask( | |
| prompt, True | |
| ): | |
| yield self.get_message(response) | |
| def for_non_stream(): | |
| return self.get_message( | |
| self.ask( | |
| prompt, | |
| False, | |
| ) | |
| ) | |
| return for_stream() if stream else for_non_stream() | |
| def get_message(self, response: dict) -> str: | |
| assert isinstance(response, dict), "Response should be of dict data-type only" | |
| if "text" in response: | |
| text = re.sub(r'\[\[\d+\]\]', '', response["text"]) | |
| return text | |
| else: | |
| return "" | |
| if __name__ == '__main__': | |
| from rich import print | |
| ai = v2() | |
| response = ai.chat("HelpingAI-9B", stream=True) | |
| for chunk in response: | |
| print(chunk, end="", flush=True) |