JetBrains AI API client

Installation

pip install jbai-client

Usage examples

Create the client:

import os

from jbai_client import JbaiPlatformClient, JbaiAuthType, JbaiEndpoint

client = JbaiPlatformClient(
    endpoint=JbaiEndpoint.STAGING,
    auth_type=JbaiAuthType.USER,
    api_key=os.getenv("JBAI_TOKEN"),
)

LLM API

Get the list of LLM profiles:

response = client.get_llm_profiles_v9()
for profile in response.profiles:
    print(profile)

Use LLM chat:

# noinspection PyTypeChecker
from jbai_client.models import ChatModelsStreamV9Request, LLMChatUserMessage, V5LLMChat

for response in client.post_llm_chat_stream_v9(
        request=ChatModelsStreamV9Request(
            prompt="test",
            profile="openai-gpt-4",
            chat=V5LLMChat(
                messages=[
                    LLMChatUserMessage(content="Tell me a joke about programmers"),
                ]
            ),
        )
):
    print(response)

Use AsyncJbaiPlatformClient for asynchronous execution:

import asyncio
import os

from jbai_client import AsyncJbaiPlatformClient, JbaiAuthType, JbaiEndpoint
from jbai_client.models import ChatModelsStreamV9Request, LLMChatUserMessage, V5LLMChat

client = AsyncJbaiPlatformClient(
    endpoint=JbaiEndpoint.STAGING,
    auth_type=JbaiAuthType.USER,
    api_key=os.getenv("JBAI_TOKEN"),
)

async def main():
    # noinspection PyTypeChecker
    async for response in client.post_llm_chat_stream_v9(
            request=ChatModelsStreamV9Request(
                prompt="test",
                profile="openai-gpt-4",
                chat=V5LLMChat(
                    messages=[
                        LLMChatUserMessage(content="Tell me a joke about programmers"),
                    ]
                ),
            )
    ):
        print(response)

asyncio.run(main())

Tasks API

Get tasks roster:

response = client.get_task_roster()
for task in response.ids:
    print(task)

Call a task (optionally specify a task tag via custom header (valid only for STAGING)):

from jbai_client import JbaiHeader
from jbai_client.models import TaskAPIStreamV5TextImproveShortenRequest

# noinspection PyTypeChecker
for response in client.text_improve_shorten_v5(
        TaskAPIStreamV5TextImproveShortenRequest(
            parameters={
                "text": "This is a simple example sentence that could be made shorter and clearer.",
                "lang": "en",
            }
        ),
        headers={
            JbaiHeader.GRAZIE_TASK_TAG: "openai-chat-gpt"
        },
):
    print(response)

P.S. Most modern IDEs support auto-import functionality, which can simplify the discovery and import of data model classes.

Data models documentation

See MODELS.md