Coverage for src/common/models/openapi_operation.py: 100%

1"""Module contains the class for an OpenAPI endpoint operation, called an  

2Operation Object in OpenAPI spec parlance.""" 

3 

4import enum 

5import uuid 

6from typing import List 

7 

8from pgvector.sqlalchemy import Vector 

9from sqlalchemy import JSON, Column, Enum, ForeignKey, Text 

10from sqlalchemy.dialects.postgresql import UUID 

11from sqlalchemy.orm import Mapped, mapped_column, relationship 

12 

13from .base import Base 

14 

15 

16# Are these the only operations we plan to support? 

17class HttpVerb(enum.Enum): 

18 GET = "GET" 

19 POST = "POST" 

20 PUT = "PUT" 

21 PATCH = "PATCH" 

22 DELETE = "DELETE" 

23 

24 

25class OpenAPIOperation(Base): # pylint: disable=too-few-public-methods 

26 """Class for an OpenAPI endpoint, called a path in OpenAPI 

27 spec parlance.""" 

28 

29 __tablename__ = "openapi_operation" 

30 

31 openapi_operation_id = Column( 

32 UUID(as_uuid=True), primary_key=True, default=uuid.uuid4 

33 ) 

34 

35 # An operation can belong to a list of servers, and is said to be defined 

36 # for each of those servers 

37 # As a constraint for CueCode, there must be exactly one server per path. 

38 openapi_server_id = Column( 

39 UUID(as_uuid=True), 

40 ForeignKey("openapi_server.openapi_server_id"), 

41 nullable=False, 

42 ) 

43 openapi_path_id = Column( 

44 UUID(as_uuid=True), ForeignKey("openapi_path.openapi_path_id"), nullable=False 

45 ) 

46 

47 http_verb: Column = Column(Enum(HttpVerb), nullable=False) 

48 

49 selection_prompt = Column(Text, nullable=False) 

50 selection_prompt_embedding = mapped_column(Vector(4096)) 

51 llm_content_gen_tool_call_spec = Column(JSON) 

52 # TODO: add embedding pylint: disable=fixme 

53 

54 path: Mapped["OpenAPIPath"] = relationship(back_populates="operations")