File size: 960 Bytes
f893655
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import requests

from langchain.llms.base import LLM


class CustomLLM(LLM):
    @property
    def _llm_type(self) -> str:
        return "custom"

    def _call(self, prompt: str, stop=None) -> str:
        r = requests.post(
            "http://localhost:8000/v1/chat/completions",
            json={
                "model": "283-vicuna-7b",
                "messages": [{"role": "user", "content": prompt}],
                "stop": stop
            },
        )
        result = r.json()
        return result["choices"][0]["message"]["content"]

    async def _acall(self, prompt: str, stop=None) -> str:
        r = requests.post(
            "http://localhost:8000/v1/chat/completions",
            json={
                "model": "283-vicuna-7b",
                "messages": [{"role": "user", "content": prompt}],
                "stop": stop
            },
        )
        result = r.json()
        return result["choices"][0]["message"]["content"]