|  | @@ -0,0 +1,87 @@
 | 
	
		
			
				|  |  | +from __future__ import annotations
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +import logging
 | 
	
		
			
				|  |  | +from typing import (
 | 
	
		
			
				|  |  | +    Any,
 | 
	
		
			
				|  |  | +    Dict,
 | 
	
		
			
				|  |  | +    List,
 | 
	
		
			
				|  |  | +    Optional,
 | 
	
		
			
				|  |  | +)
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +import requests
 | 
	
		
			
				|  |  | +from langchain.llms.utils import enforce_stop_tokens
 | 
	
		
			
				|  |  | +from pydantic import Field
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +from langchain.callbacks.manager import (
 | 
	
		
			
				|  |  | +    AsyncCallbackManagerForLLMRun,
 | 
	
		
			
				|  |  | +    CallbackManagerForLLMRun,
 | 
	
		
			
				|  |  | +)
 | 
	
		
			
				|  |  | +from langchain.llms.base import LLM
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +logger = logging.getLogger(__name__)
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +class OpenLLM(LLM):
 | 
	
		
			
				|  |  | +    """OpenLLM, supporting both in-process model
 | 
	
		
			
				|  |  | +    instance and remote OpenLLM servers.
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    If you have a OpenLLM server running, you can also use it remotely:
 | 
	
		
			
				|  |  | +        .. code-block:: python
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +            from langchain.llms import OpenLLM
 | 
	
		
			
				|  |  | +            llm = OpenLLM(server_url='http://localhost:3000')
 | 
	
		
			
				|  |  | +            llm("What is the difference between a duck and a goose?")
 | 
	
		
			
				|  |  | +    """
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    server_url: Optional[str] = None
 | 
	
		
			
				|  |  | +    """Optional server URL that currently runs a LLMServer with 'openllm start'."""
 | 
	
		
			
				|  |  | +    llm_kwargs: Dict[str, Any] = Field(default_factory=dict)
 | 
	
		
			
				|  |  | +    """Key word arguments to be passed to openllm.LLM"""
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    @property
 | 
	
		
			
				|  |  | +    def _llm_type(self) -> str:
 | 
	
		
			
				|  |  | +        return "openllm"
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    def _call(
 | 
	
		
			
				|  |  | +        self,
 | 
	
		
			
				|  |  | +        prompt: str,
 | 
	
		
			
				|  |  | +        stop: Optional[List[str]] = None,
 | 
	
		
			
				|  |  | +        run_manager: CallbackManagerForLLMRun | None = None,
 | 
	
		
			
				|  |  | +        **kwargs: Any,
 | 
	
		
			
				|  |  | +    ) -> str:
 | 
	
		
			
				|  |  | +        params = {
 | 
	
		
			
				|  |  | +            "prompt": prompt,
 | 
	
		
			
				|  |  | +            "llm_config": self.llm_kwargs
 | 
	
		
			
				|  |  | +        }
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        headers = {"Content-Type": "application/json"}
 | 
	
		
			
				|  |  | +        response = requests.post(
 | 
	
		
			
				|  |  | +            f'{self.server_url}/v1/generate',
 | 
	
		
			
				|  |  | +            headers=headers,
 | 
	
		
			
				|  |  | +            json=params
 | 
	
		
			
				|  |  | +        )
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        if not response.ok:
 | 
	
		
			
				|  |  | +            raise ValueError(f"OpenLLM HTTP {response.status_code} error: {response.text}")
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        json_response = response.json()
 | 
	
		
			
				|  |  | +        completion = json_response["responses"][0]
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        if completion:
 | 
	
		
			
				|  |  | +            completion = completion[len(prompt):]
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        if stop is not None:
 | 
	
		
			
				|  |  | +            completion = enforce_stop_tokens(completion, stop)
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        return completion
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    async def _acall(
 | 
	
		
			
				|  |  | +        self,
 | 
	
		
			
				|  |  | +        prompt: str,
 | 
	
		
			
				|  |  | +        stop: Optional[List[str]] = None,
 | 
	
		
			
				|  |  | +        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
 | 
	
		
			
				|  |  | +        **kwargs: Any,
 | 
	
		
			
				|  |  | +    ) -> str:
 | 
	
		
			
				|  |  | +        raise NotImplementedError(
 | 
	
		
			
				|  |  | +            "Async call is not supported for OpenLLM at the moment."
 | 
	
		
			
				|  |  | +        )
 |