feat: adapt proxy

Author: diogoncalves

examples/01_intro_to_llmstudio.ipynb

@@ -6,23 +6,38 @@
 
 llm = LLMProxyProvider(provider="openai", host="0.0.0.0", port="8001")
 
-result = llm.chat("What's your name", model="gpt-4o")
+result = llm.chat("Write a paragfraph about space", model="gpt-4o")
 print(result)
 
+
+response = llm.chat("Write a paragfraph about space", model="gpt-4o", is_stream=True)
+for i, chunk in enumerate(response):
+    if i%20==0:
+        print("\n")
+    if not chunk.metrics:
+        print(chunk.chat_output, end="", flush=True)
+    else:
+        print("\n\n## Metrics:")
+        print(chunk.metrics)
+
+
 import asyncio
 
 # stream
 print("\nasync stream")
 async def async_stream():
 
-    response_async = await llm.achat("What's your name", model="gpt-4o", is_stream=True)
-    async for p in response_async:
-        if "}" in p.chat_output:
-            p.chat_output
-        print("that: ",p.chat_output)
+    response_async = await llm.achat("Write a paragfraph about space", model="gpt-4o", is_stream=False)
+    print(response_async)
+    
+    response_async_stream = await llm.achat("Write a paragfraph about space", model="gpt-4o", is_stream=True)
+    async for p in response_async_stream:
+        
         # pprint(p.choices[0].delta.content==p.chat_output)
         # print("metrics: ", p.metrics)
         # print(p)
-        if p.metrics:
-            print(p)
+        if not p.metrics:
+            print(p.chat_output, end="", flush=True)
+        else:
+            print(p.metrics)
 asyncio.run(async_stream())

examples/01_intro_to_llmstudio_with_proxy.ipynb

@@ -7,7 +7,7 @@
 _engine_config = _load_engine_config()
 
 
-def LLM(provider: str, api_key: Optional[str] = None, **kwargs) -> BaseProvider:
+def LLMCore(provider: str, api_key: Optional[str] = None, **kwargs) -> BaseProvider:
     """
     Factory method to create an instance of a provider.
 
@@ -35,7 +35,7 @@ def LLM(provider: str, api_key: Optional[str] = None, **kwargs) -> BaseProvider:
     load_dotenv()
 
     def test_stuff(provider, model, api_key, **kwargs):
-        llm = LLM(provider=provider, api_key=api_key, **kwargs)
+        llm = LLMCore(provider=provider, api_key=api_key, **kwargs)
 
         latencies = {}
         chat_request = {
@@ -76,9 +76,8 @@ async def async_stream():
 
             response_async = await llm.achat(**chat_request)
             async for p in response_async:
-                if "}" in p.chat_output:
-                    p.chat_output
-                print("that: ",p.chat_output)
+                if not p.metrics:
+                    print("that: ",p.chat_output_stream)
                 # pprint(p.choices[0].delta.content==p.chat_output)
                 # print("metrics: ", p.metrics)
                 # print(p)

examples/llm_proxy.py

@@ -229,7 +229,8 @@ async def ahandle_response(
                                 if isinstance(request.chat_input, str)
                                 else request.chat_input[-1]["content"]
                             ),
-                            "chat_output": chat_output if chat_output else "",
+                            "chat_output": None,
+                            "chat_output_stream": chat_output if chat_output else "",
                             "context": (
                                 [{"role": "user", "content": request.chat_input}]
                                 if isinstance(request.chat_input, str)
@@ -276,7 +277,8 @@ async def ahandle_response(
                 if isinstance(request.chat_input, str)
                 else request.chat_input[-1]["content"]
             ),
-            "chat_output": "" if request.is_stream else output_string,
+            "chat_output": output_string,
+            "chat_output_stream": "",
             "context": (
                 [{"role": "user", "content": request.chat_input}]
                 if isinstance(request.chat_input, str)

libs/core/llmstudio_core/__init__.py

@@ -96,7 +96,13 @@ async def chat_handler(request: Request):
             provider_class = provider_registry.get(f"{provider_config.name}".lower())
             provider_instance = provider_class(provider_config)
             request_dict = await request.json()
+
             result = await provider_instance.achat(**request_dict)
+            if request_dict.get("is_stream", False):
+                async def result_generator():
+                    async for chunk in result:
+                        yield json.dumps(chunk.dict())
+                return StreamingResponse(result_generator(), media_type="application/json")
             return result
 
         return chat_handler

libs/core/llmstudio_core/providers/provider.py

@@ -1,27 +1,40 @@
 import asyncio
-from typing import Any, Coroutine, Dict, List, Union
+import json
+from typing import Any, Coroutine, Dict, List, Optional, Union
 
+from pydantic import BaseModel
 import requests
-from libs.core.llmstudio_core.providers.provider import BaseProvider, ProviderABC
+from llmstudio_core.providers.provider import ProviderABC
 from llmstudio.server import is_server_running
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from tqdm.asyncio import tqdm_asyncio
 
-from llmstudio.config import ENGINE_HOST, ENGINE_PORT
 from llmstudio.llm.semaphore import DynamicSemaphore
 
+
+class ProxyConfig(BaseModel):
+    host: Optional[str] = None
+    port: Optional[str] = None
+    url: Optional[str] = None
+    username: Optional[str] = None
+    password: Optional[str] = None
+    def __init__(self, **data):
+        super().__init__(**data)
+        if (self.host is None and self.port is None) and self.url is None:
+            raise ValueError("Either both 'host' and 'port' must be provided, or 'url' must be specified.")
+        
+        
 class LLMProxyProvider(ProviderABC):
     def __init__(self, provider: str,                 
-                 host: str,
-                 port: str,
-        **kwargs):
+                 proxy_config: ProxyConfig):
         self.provider = provider
-        self.engine_host = host
-        self.engine_port = port
-        if is_server_running(host=host, port=port):
-            print(f"Connected to LLMStudio Proxy @ {host}:{port}")
+        
+        self.engine_host = proxy_config.host
+        self.engine_port = proxy_config.port
+        if is_server_running(host=self.engine_host, port=self.engine_port):
+            print(f"Connected to LLMStudio Proxy @ {self.engine_host}:{self.engine_port}")
         else:
-            raise Exception(f"LLMStudio Proxy is not running @ {host}:{port}")
+            raise Exception(f"LLMStudio Proxy is not running @ {self.engine_host}:{self.engine_port}")
 
     @staticmethod
     def _provider_config_name():
@@ -59,7 +72,7 @@ def chat(self, chat_input: str,
     def generate_chat(self, response):
         for chunk in response.iter_content(chunk_size=None):
             if chunk:
-                yield chunk.decode("utf-8")
+                yield ChatCompletionChunk(**json.loads(chunk.decode("utf-8")))
 
     async def achat(self, chat_input: Any, 
               model: str, 
@@ -233,4 +246,4 @@ async def async_stream(self, model:str, chat_input: str, retries: int, parameter
 
         for chunk in response.iter_content(chunk_size=None):
             if chunk:
-                yield chunk.decode("utf-8")
+                yield ChatCompletionChunk(**json.loads(chunk.decode("utf-8")))

llmstudio/engine/__init__.py

@@ -1,43 +1,44 @@
 from typing import Any, Coroutine, Optional
-from llmstudio_core import LLM as LLM_factory
+from llmstudio_core import LLMCore
 from llmstudio_core.providers.provider import ProviderABC
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from pydantic import BaseModel
 
-from llmstudio.engine.provider import LLMProxyProvider
+from llmstudio.engine.provider import LLMProxyProvider, ProxyConfig
 from llmstudio.tracking.database import create_tracking_engine
 from llmstudio.tracking.logs import crud, schemas
 
-from sqlalchemy.orm import declarative_base, sessionmaker
-
-
-class ProxyConfig(BaseModel):
-    host: str
-    port: int
-    username: Optional[str] = None
-    password: Optional[str] = None
+from sqlalchemy.orm import sessionmaker
 
 class TrackingConfig(BaseModel):
-    database_uri: str
+    database_uri: Optional[str] = None
+    host: Optional[str] = None
+    port: Optional[int] = None
+    url: Optional[str] = None
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        if (self.host and self.port) or self.url or self.database_uri:
+            raise ValueError("You must provide either both 'host' and 'port', or 'url', or 'database_uri'.")
 
 
 class LLM(ProviderABC):
 
 
     def __init__(self,
                  provider: str,
+                 api_key: Optional[str] = None,
                  proxy_config: Optional[ProxyConfig] = None,
                  tracking_config: Optional[TrackingConfig] = None,
                  **kwargs):
 
         if proxy_config is not None:
             self._provider = LLMProxyProvider(provider=provider,
-                                              host=proxy_config.host,
-                                              port=proxy_config.port,
-                                              **kwargs
-                                              )
+                                              proxy_config=proxy_config)
         else:
-            self._provider = LLM_factory(provider, **kwargs)
+            self._provider = LLMCore(provider=provider,
+                                     api_key=api_key, 
+                                     **kwargs)
 
         self._session_local = None
         if tracking_config is not None: