diff --git a/main.py b/main.py index 04a78e7..8e1f356 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,6 @@ from prank_line_crafter.story_generator import generate_complete_story from prank_line_crafter.sender import send_story -import os + def get_names_from_file(filename): """Read names from the specified file.""" @@ -12,19 +12,21 @@ def get_names_from_file(filename): print(f"The file {filename} does not exist.") return [] + def main(): names_file = 'names.txt' names = get_names_from_file(names_file) for i, name in enumerate(names, start=1): print(f"{i}: Processing {name}... ", end="") - + # Generate a funny story for the current name story = generate_complete_story(name) - + # Send the story to the server result = send_story(name, story) print(result) + if __name__ == "__main__": main() diff --git a/prank_line_crafter/config.py b/prank_line_crafter/config.py index 33836a6..e8d806f 100644 --- a/prank_line_crafter/config.py +++ b/prank_line_crafter/config.py @@ -7,4 +7,4 @@ URL = os.getenv("URL") if not LLAMA_API_TOKEN or not URL: - raise ValueError("Missing required environment variables: LLAMA_API_TOKEN or URL") + raise ValueError("Missing: LLAMA_API_TOKEN or URL") diff --git a/prank_line_crafter/sender.py b/prank_line_crafter/sender.py index 1e3a440..f440043 100644 --- a/prank_line_crafter/sender.py +++ b/prank_line_crafter/sender.py @@ -1,6 +1,7 @@ import requests from .config import URL + def send_story(name, story): """Send the generated story to the server.""" data = { diff --git a/prank_line_crafter/story_generator.py b/prank_line_crafter/story_generator.py index febfd97..ea8058c 100644 --- a/prank_line_crafter/story_generator.py +++ b/prank_line_crafter/story_generator.py @@ -1,12 +1,14 @@ from llamaapi import LlamaAPI from .config import LLAMA_API_TOKEN + def generate_complete_story(name): """Generate a funny story about the given name using LlamaAPI.""" llama = LlamaAPI(LLAMA_API_TOKEN) api_request_json = { "messages": [ - {"role": "user", "content": f"Generate a funny story about a person named {name}."} + {"role": "user", + "content": f"Generate a funny story about a person named {name}"} ], "max_tokens": 500, "stream": False diff --git a/tests/test_main.py b/tests/test_main.py index c637baf..4929d93 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,10 +1,9 @@ import unittest from unittest.mock import patch, MagicMock from prank_line_crafter.story_generator import generate_complete_story -from main import get_names_from_file, main +from main import get_names_from_file import os -import io -import sys + class TestMain(unittest.TestCase): @@ -40,5 +39,6 @@ def test_get_names_from_file(self): # Assert that the names list matches the expected list self.assertEqual(names, ['Alice', 'Bob', 'Charlie']) + if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main()