API-Handler commited on
Commit
4e2263c
1 Parent(s): b95d0b7

Upload 5 files

Browse files
Files changed (5) hide show
  1. Dockerfile +20 -0
  2. deepinfra_handler.py +65 -0
  3. inference.py +133 -0
  4. main.py +53 -0
  5. requirements.txt +4 -0
Dockerfile ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Use an official Python runtime as the base image
2
+ FROM python:3.9-slim
3
+
4
+ # Set the working directory in the container
5
+ WORKDIR /app
6
+
7
+ # Copy the requirements file into the container
8
+ COPY requirements.txt .
9
+
10
+ # Install the required packages
11
+ RUN pip install --no-cache-dir -r requirements.txt
12
+
13
+ # Copy the rest of the application code into the container
14
+ COPY . .
15
+
16
+ # Expose the port that FastAPI will run on
17
+ EXPOSE 7860
18
+
19
+ # Command to run the FastAPI application
20
+ CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
deepinfra_handler.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import requests
3
+ from typing import Dict, Any, Generator, Optional
4
+
5
+ class DeepInfraHandler:
6
+ API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
7
+
8
+ def __init__(self):
9
+ self.headers = {
10
+ "Accept": "text/event-stream",
11
+ "Accept-Encoding": "gzip, deflate, br, zstd",
12
+ "Content-Type": "application/json",
13
+ "Connection": "keep-alive",
14
+ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36",
15
+ }
16
+
17
+ def _prepare_payload(self, **kwargs) -> Dict[str, Any]:
18
+ """Prepare the payload for the API request"""
19
+ return {
20
+ "model": kwargs.get("model"),
21
+ "messages": kwargs.get("messages"),
22
+ "temperature": kwargs.get("temperature", 0.7),
23
+ "max_tokens": kwargs.get("max_tokens", 4096),
24
+ "top_p": kwargs.get("top_p", 1.0),
25
+ "frequency_penalty": kwargs.get("frequency_penalty", 0.0),
26
+ "presence_penalty": kwargs.get("presence_penalty", 0.0),
27
+ "stop": kwargs.get("stop", []),
28
+ "stream": kwargs.get("stream", False)
29
+ }
30
+
31
+ def generate_completion(self, **kwargs) -> Any:
32
+ """Generate completion based on streaming preference"""
33
+ payload = self._prepare_payload(**kwargs)
34
+
35
+ response = requests.post(
36
+ self.API_URL,
37
+ headers=self.headers,
38
+ json=payload,
39
+ stream=payload["stream"]
40
+ )
41
+
42
+ if payload["stream"]:
43
+ return self._handle_streaming_response(response)
44
+ return self._handle_regular_response(response)
45
+
46
+ def _handle_streaming_response(self, response) -> Generator[str, None, None]:
47
+ """Handle streaming response from the API"""
48
+ for line in response.iter_lines(decode_unicode=True):
49
+ if line.startswith("data:"):
50
+ try:
51
+ content = json.loads(line[5:])
52
+ if content == "[DONE]":
53
+ continue
54
+ delta_content = content.get("choices", [{}])[0].get("delta", {}).get("content")
55
+ if delta_content:
56
+ yield delta_content
57
+ except:
58
+ continue
59
+
60
+ def _handle_regular_response(self, response) -> Dict[str, Any]:
61
+ """Handle regular (non-streaming) response from the API"""
62
+ try:
63
+ return response.json()
64
+ except Exception as e:
65
+ raise Exception(f"Error processing response: {str(e)}")
inference.py ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import requests
2
+ import json
3
+ from typing import Union, Dict, Generator
4
+ import time
5
+
6
+ class ChatCompletionTester:
7
+ def __init__(self, base_url: str = "http://localhost:8000"):
8
+ self.base_url = base_url
9
+ self.endpoint = f"{base_url}/chat/completions"
10
+
11
+ def create_test_payload(self, stream: bool = False) -> Dict:
12
+ """Create a sample payload for testing"""
13
+ return {
14
+ "model": "mistralai/Mixtral-8x22B-Instruct-v0.1",
15
+ "messages": [
16
+ {"role": "system", "content": "You are a helpful assistant."},
17
+ {"role": "user", "content": "What is the capital of France?"}
18
+ ],
19
+ "temperature": 0.7,
20
+ "max_tokens": 4096,
21
+ "stream": stream
22
+ }
23
+
24
+ def test_non_streaming(self) -> Union[Dict, None]:
25
+ """Test non-streaming response"""
26
+ print("\n=== Testing Non-Streaming Response ===")
27
+ try:
28
+ payload = self.create_test_payload(stream=False)
29
+ print("Sending request...")
30
+
31
+ response = requests.post(
32
+ self.endpoint,
33
+ json=payload,
34
+ headers={"Content-Type": "application/json"}
35
+ )
36
+
37
+ if response.status_code == 200:
38
+ result = response.json()
39
+ content = result['choices'][0]['message']['content']
40
+ print("\nResponse received successfully!")
41
+ print(f"Content: {content}")
42
+ return result
43
+ else:
44
+ print(f"Error: Status code {response.status_code}")
45
+ print(f"Response: {response.text}")
46
+ return None
47
+
48
+ except Exception as e:
49
+ print(f"Error during non-streaming test: {str(e)}")
50
+ return None
51
+
52
+ def test_streaming(self) -> Union[str, None]:
53
+ """Test streaming response"""
54
+ print("\n=== Testing Streaming Response ===")
55
+ try:
56
+ payload = self.create_test_payload(stream=True)
57
+ print("Sending request...")
58
+
59
+ response = requests.post(
60
+ self.endpoint,
61
+ json=payload,
62
+ headers={"Content-Type": "application/json"},
63
+ stream=True
64
+ )
65
+
66
+ if response.status_code == 200:
67
+ print("\nReceiving streaming response:")
68
+ full_response = ""
69
+ for line in response.iter_lines(decode_unicode=True):
70
+ if line:
71
+ if line.startswith("data: "):
72
+ try:
73
+ data = json.loads(line[6:])
74
+ if data == "[DONE]":
75
+ continue
76
+ content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
77
+ if content:
78
+ print(content, end="", flush=True)
79
+ full_response += content
80
+ except json.JSONDecodeError:
81
+ continue
82
+ print("\n\nStreaming completed!")
83
+ return full_response
84
+ else:
85
+ print(f"Error: Status code {response.status_code}")
86
+ print(f"Response: {response.text}")
87
+ return None
88
+
89
+ except Exception as e:
90
+ print(f"Error during streaming test: {str(e)}")
91
+ return None
92
+
93
+ def run_all_tests(self):
94
+ """Run both streaming and non-streaming tests"""
95
+ print("Starting API endpoint tests...")
96
+
97
+ # Test server connectivity
98
+ try:
99
+ requests.get(self.base_url)
100
+ print("✓ Server is accessible")
101
+ except requests.exceptions.ConnectionError:
102
+ print("✗ Server is not accessible. Please ensure the FastAPI server is running.")
103
+ return
104
+
105
+ # Run tests with timing
106
+ start_time = time.time()
107
+
108
+ # Test non-streaming
109
+ non_streaming_result = self.test_non_streaming()
110
+ if non_streaming_result:
111
+ print("✓ Non-streaming test passed")
112
+ else:
113
+ print("✗ Non-streaming test failed")
114
+
115
+ # Test streaming
116
+ streaming_result = self.test_streaming()
117
+ if streaming_result:
118
+ print("✓ Streaming test passed")
119
+ else:
120
+ print("✗ Streaming test failed")
121
+
122
+ end_time = time.time()
123
+ print(f"\nAll tests completed in {end_time - start_time:.2f} seconds")
124
+
125
+ def main():
126
+ # Create tester instance
127
+ tester = ChatCompletionTester()
128
+
129
+ # Run all tests
130
+ tester.run_all_tests()
131
+
132
+ if __name__ == "__main__":
133
+ main()
main.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException
2
+ from fastapi.responses import StreamingResponse
3
+ from pydantic import BaseModel, Field
4
+ from typing import List, Optional, Any, Dict
5
+ from deepinfra_handler import DeepInfraHandler
6
+ import json
7
+
8
+ app = FastAPI()
9
+ api_handler = DeepInfraHandler()
10
+
11
+ class Message(BaseModel):
12
+ role: str
13
+ content: str
14
+
15
+ class ChatCompletionRequest(BaseModel):
16
+ model: str
17
+ messages: List[Message]
18
+ temperature: Optional[float] = Field(default=0.7, ge=0.0, le=2.0)
19
+ max_tokens: Optional[int] = Field(default=4096, ge=1)
20
+ top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
21
+ frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
22
+ presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
23
+ stop: Optional[List[str]] = Field(default=[])
24
+ stream: Optional[bool] = Field(default=False)
25
+
26
+ @app.post("/chat/completions")
27
+ async def chat_completions(request: ChatCompletionRequest):
28
+ try:
29
+ # Convert request to dictionary
30
+ params = request.dict()
31
+
32
+ if request.stream:
33
+ # Handle streaming response
34
+ def generate():
35
+ for chunk in api_handler.generate_completion(**params):
36
+ yield f"data: {json.dumps({'choices': [{'delta': {'content': chunk}}]})}\n\n"
37
+ yield "data: [DONE]\n\n"
38
+
39
+ return StreamingResponse(
40
+ generate(),
41
+ media_type="text/event-stream"
42
+ )
43
+
44
+ # Handle regular response
45
+ response = api_handler.generate_completion(**params)
46
+ return response
47
+
48
+ except Exception as e:
49
+ raise HTTPException(status_code=500, detail=str(e))
50
+
51
+ if __name__ == "__main__":
52
+ import uvicorn
53
+ uvicorn.run(app, host="0.0.0.0", port=8000)
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ fastapi>=0.104.1
2
+ uvicorn>=0.24.0
3
+ requests>=2.31.0
4
+ pydantic>=2.5.2