1  # Usage:
 2  #! ./llama-server -m some-model.gguf &
 3  #! pip install pydantic
 4  #! python json-schema-pydantic-example.py
 5  
 6  from pydantic import BaseModel, TypeAdapter
 7  from annotated_types import MinLen
 8  from typing import Annotated, List, Optional
 9  import json, requests
10  
11  if True:
12  
13      def create_completion(*, response_model=None, endpoint="http://localhost:8080/v1/chat/completions", messages, **kwargs):
14          '''
15          Creates a chat completion using an OpenAI-compatible endpoint w/ JSON schema support
16          (llama.cpp server, llama-cpp-python, Anyscale / Together...)
17  
18          The response_model param takes a type (+ supports Pydantic) and behaves just as w/ Instructor (see below)
19          '''
20          if response_model:
21              type_adapter = TypeAdapter(response_model)
22              schema = type_adapter.json_schema()
23              messages = [{
24                  "role": "system",
25                  "content": f"You respond in JSON format with the following schema: {json.dumps(schema, indent=2)}"
26              }] + messages
27              response_format={"type": "json_object", "schema": schema}
28  
29          data = requests.post(endpoint, headers={"Content-Type": "application/json"},
30                               json=dict(messages=messages, response_format=response_format, **kwargs)).json()
31          if 'error' in data:
32              raise Exception(data['error']['message'])
33  
34          content = data["choices"][0]["message"]["content"]
35          return type_adapter.validate_json(content) if type_adapter else content
36  
37  else:
38  
39      # This alternative branch uses Instructor + OpenAI client lib.
40      # Instructor support streamed iterable responses, retry & more.
41      # (see https://python.useinstructor.com/)
42      #! pip install instructor openai
43      import instructor, openai
44      client = instructor.patch(
45          openai.OpenAI(api_key="123", base_url="http://localhost:8080"),
46          mode=instructor.Mode.JSON_SCHEMA)
47      create_completion = client.chat.completions.create
48  
49  
50  if __name__ == '__main__':
51  
52      class QAPair(BaseModel):
53          question: str
54          concise_answer: str
55          justification: str
56  
57      class PyramidalSummary(BaseModel):
58          title: str
59          summary: str
60          question_answers: Annotated[List[QAPair], MinLen(2)]
61          sub_sections: Optional[Annotated[List['PyramidalSummary'], MinLen(2)]]
62  
63      print("# Summary\n", create_completion(
64          model="...",
65          response_model=PyramidalSummary,
66          messages=[{
67              "role": "user",
68              "content": f"""
69                  You are a highly efficient corporate document summarizer.
70                  Create a pyramidal summary of an imaginary internal document about our company processes
71                  (starting high-level, going down to each sub sections).
72                  Keep questions short, and answers even shorter (trivia / quizz style).
73              """
74          }]))