Custom LLM wrapper: any API

from llama_index.core.llms import LLM, CompletionResponse, ChatResponse, ChatMessage
from llama_index.core.base.llms.types import LLMMetadata
import httpx
import json

class CustomAPILLM(LLM):
    """Wrapper for a custom inference API."""

    model: str = "custom-model-v1"
    api_base: str = "https://api.example.com"
    api_key: str = ""
    temperature: float = 0.7
    max_tokens: int = 512

    @property
    def metadata(self) -> LLMMetadata:
        """Return metadata about the LLM."""
        return LLMMetadata(
            name=self.__class__.__name__,
            model_name=self.model,
            context_window=4096,
            num_output=self.max_tokens,
        )

    def complete(
        self,
        prompt: str,
        **kwargs,
    ) -> CompletionResponse:
        """Generate text completion from a single prompt."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": self.model,
            "prompt": prompt,
            "temperature": kwargs.get("temperature", self.temperature),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
        }

        with httpx.Client() as client:
            response = client.post(
                f"{self.api_base}/v1/completions",
                json=payload,
                headers=headers,
                timeout=30.0,
            )
            response.raise_for_status()
            data = response.json()

        text = data["choices"][0]["text"]
        return CompletionResponse(
            text=text,
            raw=data,
            additional_kwargs={"finish_reason": data["choices"][0].get("finish_reason")},
        )

    def chat(
        self,
        messages: list[ChatMessage],
        **kwargs,
    ) -> ChatResponse:
        """Generate a chat response from a list of messages."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        formatted_messages = [
            {
                "role": message.role,
                "content": message.content,
            }
            for message in messages
        ]

        payload = {
            "model": self.model,
            "messages": formatted_messages,
            "temperature": kwargs.get("temperature", self.temperature),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
        }

        with httpx.Client() as client:
            response = client.post(
                f"{self.api_base}/v1/chat/completions",
                json=payload,
                headers=headers,
                timeout=30.0,
            )
            response.raise_for_status()
            data = response.json()

        message_content = data["choices"][0]["message"]["content"]
        return ChatResponse(
            message=ChatMessage(role="assistant", content=message_content),
            raw=data,
            additional_kwargs={"finish_reason": data["choices"][0].get("finish_reason")},
        )


from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.vector_stores.faiss import FaissVectorStore
from faiss import IndexFlatL2
import tempfile
import os

llm = CustomAPILLM(
    api_base="https://api.example.com",
    api_key="your-api-key-here",
    model="custom-model-v1",
    temperature=0.5,
    max_tokens=256,
)

Settings.llm = llm

with tempfile.TemporaryDirectory() as tmpdir:
    sample_file = os.path.join(tmpdir, "sample.txt")
    with open(sample_file, "w") as f:
        f.write("Artificial intelligence is transforming how we work and live.")

    documents = SimpleDirectoryReader(tmpdir).load_data()

    faiss_index = IndexFlatL2(384)
    vector_store = FaissVectorStore(faiss_index=faiss_index)
    index = VectorStoreIndex.from_documents(
        documents,
        vector_store=vector_store,
    )

    query_engine = index.as_query_engine()
    result = query_engine.query("What is artificial intelligence?")
    print("Query Result:")
    print(result)