Get Python code closer to PEP 8 standard

.gitignore

@@ -46,3 +46,4 @@ last_recording.mp4
 next-env.d.ts
 
 !samples/chatbot/server/genkit-ai-vertexai-*.tgz
+.idea/

py/packages/genkit/src/genkit/ai/model.py

@@ -2,6 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from typing import Callable
-from genkit.core.types import GenerateRequest, GenerateResponse
+from genkit.core.schemas import GenerateRequest, GenerateResponse
 
 ModelFn = Callable[[GenerateRequest], GenerateResponse]

py/packages/genkit/src/genkit/ai/prompt.py

@@ -2,7 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-from genkit.core.types import GenerateRequest
 from typing import Callable, Optional, Any
+from genkit.core.schemas import GenerateRequest
+
 
 PromptFn = Callable[[Optional[Any]], GenerateRequest]

py/packages/genkit/src/genkit/core/action.py

@@ -1,12 +1,9 @@
 # Copyright 2025 Google LLC
-# SPDX-License-Identifier: Apache-2.0
-
-
+# SPDX-License-Identifier: Apache-2.
 import inspect
 import json
 
 from typing import Dict, Optional, Callable, Any
-
 from pydantic import ConfigDict, BaseModel, TypeAdapter
 
 from genkit.core.tracing import tracer
@@ -20,29 +17,32 @@ class ActionResponse(BaseModel):
 
 
 class Action:
+    INPUT_KEY = 'inputSchema'
+    OUTPUT_KEY = 'outputSchema'
+    RETURN = 'return'
+
     def __init__(
         self,
-        type: str,
+        action_type: str,
         name: str,
         fn: Callable,
-        description: str | None = None,
+        description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
-        spanMetadata: Optional[Dict[str, str]] = None,
+        span_metadata: Optional[Dict[str, str]] = None,
     ):
-        self.type = type
+        # TODO(Tatsiana Havina): separate a long constructor into methods.
+        self.type = action_type
         self.name = name
 
-        def fnToCall(*args, **kwargs):
+        def fn_to_call(*args, **kwargs):
             with tracer.start_as_current_span(name) as span:
-                traceId = str(span.get_span_context().trace_id)
-                span.set_attribute('genkit:type', type)
+                trace_id = str(span.get_span_context().trace_id)
+                span.set_attribute('genkit:type', action_type)
                 span.set_attribute('genkit:name', name)
 
-                if spanMetadata is not None:
-                    for spanMetaKey in spanMetadata:
-                        span.set_attribute(
-                            spanMetaKey, spanMetadata[spanMetaKey]
-                        )
+                if span_metadata is not None:
+                    for meta_key in span_metadata:
+                        span.set_attribute(meta_key, span_metadata[meta_key])
 
                 if len(args) > 0:
                     if isinstance(args[0], BaseModel):
@@ -63,33 +63,30 @@ def fnToCall(*args, **kwargs):
                 else:
                     span.set_attribute('genkit:output', json.dumps(output))
 
-                return ActionResponse(response=output, traceId=traceId)
+                return ActionResponse(response=output, traceId=trace_id)
 
-        self.fn = fnToCall
+        self.fn = fn_to_call
         self.description = description
-        self.metadata = metadata
-        if self.metadata is None:
-            self.metadata = {}
-
-        inputSpec = inspect.getfullargspec(fn)
-        actionArgs = list(
-            filter(lambda k: k != 'return', inputSpec.annotations)
-        )
-        if len(actionArgs) > 1:
+        self.metadata = metadata if metadata else {}
+
+        input_spec = inspect.getfullargspec(fn)
+        action_args = [k for k in input_spec.annotations if k != self.RETURN]
+
+        if len(action_args) > 1:
             raise Exception('can only have one arg')
-        if len(actionArgs) > 0:
-            ta = TypeAdapter(inputSpec.annotations[actionArgs[0]])
-            self.inputSchema = ta.json_schema()
-            self.inputType = ta
-            self.metadata['inputSchema'] = self.inputSchema
+        if len(action_args) > 0:
+            type_adapter = TypeAdapter(input_spec.annotations[action_args[0]])
+            self.input_schema = type_adapter.json_schema()
+            self.input_type = type_adapter
+            self.metadata[self.INPUT_KEY] = self.input_schema
         else:
-            self.inputSchema = TypeAdapter(Any).json_schema()
-            self.metadata['inputSchema'] = self.inputSchema
+            self.input_schema = TypeAdapter(Any).json_schema()
+            self.metadata[self.INPUT_KEY] = self.input_schema
 
-        if 'return' in inputSpec.annotations:
-            ta = TypeAdapter(inputSpec.annotations['return'])
-            self.outputSchema = ta.json_schema()
-            self.metadata['outputSchema'] = self.outputSchema
+        if self.RETURN in input_spec.annotations:
+            type_adapter = TypeAdapter(input_spec.annotations[self.RETURN])
+            self.output_schema = type_adapter.json_schema()
+            self.metadata[self.OUTPUT_KEY] = self.output_schema
         else:
-            self.outputSchema = TypeAdapter(Any).json_schema()
-            self.metadata['outputSchema'] = self.outputSchema
+            self.output_schema = TypeAdapter(Any).json_schema()
+            self.metadata[self.OUTPUT_KEY] = self.output_schema

py/packages/genkit/src/genkit/core/reflection.py

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+"""Exposes an API for inspecting and interacting with Genkit in development."""
+
 import json
 
 from http.server import BaseHTTPRequestHandler
@@ -10,9 +12,16 @@
 from genkit.core.registry import Registry
 
 
-def MakeReflectionServer(registry: Registry):
+def make_reflection_server(registry: Registry):
+    """Returns a ReflectionServer class."""
+
     class ReflectionServer(BaseHTTPRequestHandler):
+        """Exposes an API for local development."""
+
+        ENCODING = 'utf-8'
+
         def do_GET(self):
+            """Handles GET requests."""
             if self.path == '/api/__health':
                 self.send_response(200)
 
@@ -22,10 +31,10 @@ def do_GET(self):
                 self.end_headers()
 
                 actions = {}
-                for type in registry.actions:
-                    for name in registry.actions[type]:
-                        action = registry.lookup_action(type, name)
-                        key = f'/{type}/{name}'
+                for action_type in registry.actions:
+                    for name in registry.actions[action_type]:
+                        action = registry.lookup_action(action_type, name)
+                        key = f'/{action_type}/{name}'
                         actions[key] = {
                             'key': key,
                             'name': action.name,
@@ -34,31 +43,34 @@ def do_GET(self):
                             'metadata': action.metadata,
                         }
 
-                self.wfile.write(bytes(json.dumps(actions), 'utf-8'))
+                self.wfile.write(bytes(json.dumps(actions), self.ENCODING))
 
             else:
                 self.send_response(404)
                 self.end_headers()
 
         def do_POST(self):
+            """Handles POST requests."""
             if self.path == '/api/notify':
                 self.send_response(200)
                 self.end_headers()
 
             elif self.path == '/api/runAction':
                 content_len = int(self.headers.get('Content-Length'))
                 post_body = self.rfile.read(content_len)
-                payload = json.loads(post_body.decode(encoding='utf-8'))
+                payload = json.loads(post_body.decode(encoding=self.ENCODING))
                 print(payload)
                 action = registry.lookup_by_absolute_name(payload['key'])
                 if '/flow/' in payload['key']:
-                    input = action.inputType.validate_python(
+                    input_action = action.inputType.validate_python(
                         payload['input']['start']['input']
                     )
                 else:
-                    input = action.inputType.validate_python(payload['input'])
+                    input_action = action.inputType.validate_python(
+                        payload['input']
+                    )
 
-                output = action.fn(input)
+                output = action.fn(input_action)
 
                 self.send_response(200)
                 self.send_header('x-genkit-version', '0.9.1')
@@ -73,7 +85,7 @@ def do_POST(self):
                             + ', "traceId": "'
                             + output.traceId
                             + '"}',
-                            'utf-8',
+                            self.ENCODING,
                         )
                     )
                 else:
@@ -85,7 +97,7 @@ def do_POST(self):
                                     'telemetry': {'traceId': output.traceId},
                                 }
                             ),
-                            'utf-8',
+                            self.ENCODING,
                         )
                     )
 

py/packages/genkit/src/genkit/core/registry.py

@@ -2,21 +2,25 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-from genkit.core.action import Action
+"""The registry is used to store and lookup resources."""
+
 from typing import Dict
+from genkit.core.action import Action
 
 
 class Registry:
+    """Stores actions, trace stores, flow state stores, plugins, and schemas."""
+
     actions: Dict[str, Dict[str, Action]] = {}
 
-    def register_action(self, type: str, name: str, action: Action):
-        if type not in self.actions:
-            self.actions[type] = {}
-        self.actions[type][name] = action
+    def register_action(self, action_type: str, name: str, action: Action):
+        if action_type not in self.actions:
+            self.actions[action_type] = {}
+        self.actions[action_type][name] = action
 
-    def lookup_action(self, type: str, name: str):
-        if type in self.actions and name in self.actions[type]:
-            return self.actions[type][name]
+    def lookup_action(self, action_type: str, name: str):
+        if action_type in self.actions and name in self.actions[action_type]:
+            return self.actions[action_type][name]
         return None
 
     def lookup_by_absolute_name(self, name: str):

py/packages/genkit/src/genkit/core/tracing.py

@@ -2,13 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+"""Collects telemetry."""
+
 import json
 import os
-import requests
 import sys
-
 from typing import Any, Dict, Sequence
-from opentelemetry import trace
+
+import requests
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import (
     SpanExporter,
@@ -29,8 +30,8 @@ class TelemetryServerSpanExporter(SpanExporter):
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         for span in spans:
-            spanData = {'traceId': f'{span.context.trace_id}', 'spans': {}}
-            spanData['spans'][span.context.span_id] = {
+            span_data = {'traceId': f'{span.context.trace_id}', 'spans': {}}
+            span_data['spans'][span.context.span_id] = {
                 'spanId': f'{span.context.span_id}',
                 'traceId': f'{span.context.trace_id}',
                 'startTime': span.start_time / 1000000,
@@ -40,13 +41,13 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
                 # "links": span.links,
                 'spanKind': trace_api.SpanKind(span.kind).name,
                 'parentSpanId': f'{span.parent.span_id}'
-                if span.parent is not None
+                if span.parent
                 else None,
                 'status': {
                     'code': trace_api.StatusCode(span.status.status_code).value,
                     'description': span.status.description,
                 }
-                if span.status is not None
+                if span.status
                 else None,
                 'instrumentationLibrary': {
                     'name': 'genkit-tracer',
@@ -62,18 +63,18 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
                 #     })),
                 # },
             }
-            if spanData['spans'][span.context.span_id]['parentSpanId'] is None:
-                del spanData['spans'][span.context.span_id]['parentSpanId']
+            if not span_data['spans'][span.context.span_id]['parentSpanId']:
+                del span_data['spans'][span.context.span_id]['parentSpanId']
 
-            if span.parent is None:
-                spanData['displayName'] = span.name
-                spanData['startTime'] = span.start_time
-                spanData['endTime'] = span.end_time
+            if not span.parent:
+                span_data['displayName'] = span.name
+                span_data['startTime'] = span.start_time
+                span_data['endTime'] = span.end_time
 
             # TODO: telemetry server URL must be dynamic, whatever tools notification says
             requests.post(
                 'http://localhost:4033/api/traces',
-                data=json.dumps(spanData),
+                data=json.dumps(span_data),
                 headers={
                     'Content-Type': 'application/json',
                     'Accept': 'application/json',
@@ -99,7 +100,7 @@ def convert_attributes(attributes: Dict[str, Any]) -> Dict[str, Any]:
     processor = SimpleSpanProcessor(TelemetryServerSpanExporter())
     provider.add_span_processor(processor)
     # Sets the global default tracer provider
-    trace.set_tracer_provider(provider)
-    tracer = trace.get_tracer('genkit-tracer', 'v1', provider)
+    trace_api.set_tracer_provider(provider)
+    tracer = trace_api.get_tracer('genkit-tracer', 'v1', provider)
 else:
-    tracer = trace.get_tracer('genkit-tracer', 'v1')
+    tracer = trace_api.get_tracer('genkit-tracer', 'v1')